diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c7662c713..4cbcb4d08 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,5 +1,5 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.14.0 +ARG VLLM_BASE_TAG=v0.15.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} @@ -12,8 +12,8 @@ RUN apt-get update && \ # Install vLLM wheel for the specific commit # URL format: https://wheels.vllm.ai// -# Current vLLM commit: 706f123b23b75097da9a3f51ce538bd784fdc75b -RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/706f123b23b75097da9a3f51ce538bd784fdc75b +# Current vLLM commit: f176443446f659dbab5315e056e605d8984fd976 (v0.15.0) +RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/f176443446f659dbab5315e056e605d8984fd976 # Install vllm-omni into the same uv-managed Python environment used by the base image. RUN uv pip install --system --no-cache-dir ".[dev]" diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index d0d5a0723..50e7b8dbc 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -59,10 +59,21 @@ class OmniGenerationScheduler(VLLMScheduler): # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + # OMNI: Track requests that are already finished (e.g., marked by connector) + # These should be removed from running and not scheduled + already_finished_reqs: set[Request] = set() while req_index < len(self.running) and token_budget > 0: request = self.running[req_index] if self.omni_connector is not None: get_chunk_for_generation(self.omni_connector, request) + + # OMNI: Skip requests that are already finished or not in self.requests + # This can happen when connector marks request as finished + if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests: + already_finished_reqs.add(request) + req_index += 1 + continue + num_computed_tokens = request.num_computed_tokens required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1) num_new_tokens = min(required_tokens, token_budget) @@ -86,12 +97,24 @@ class OmniGenerationScheduler(VLLMScheduler): scheduled_running_reqs.append(request) req_index += 1 + # OMNI: Remove already finished requests from running queue + if already_finished_reqs: + self.running = remove_all(self.running, already_finished_reqs) + # Fast path selection and scheduling (treat all as diffusion requests, # independent of pooling_params) while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: request = self.waiting.peek_request() if self.omni_connector is not None: get_chunk_for_generation(self.omni_connector, request) + + # OMNI: Skip requests that are already finished or not in self.requests + # This can happen when connector marks request as finished + if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests: + # Pop the finished request from waiting queue and don't schedule it + self.waiting.pop_request() + continue + # Uniformly treat as diffusion. A feature flag can be added later # via config or request tag. diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ec2a18ac4..f3623ee27 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -194,11 +194,10 @@ class AsyncOmni(OmniBase): if stage.vllm_config is not None and stage.tokenizer is not None: try: vllm_config = stage.vllm_config - tokenizer = stage.tokenizer # Initialize input_processor + # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config self.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=tokenizer, ) # Initialize model_config self.model_config = vllm_config.model_config @@ -673,6 +672,17 @@ class AsyncOmni(OmniBase): return stage.is_tracing_enabled return False + @property + def renderer(self): + """Return the renderer from input_processor if available. + + OMNI: Required by upstream OpenAIServingModels.__init__ which + accesses engine_client.renderer. + """ + if self.input_processor is not None: + return self.input_processor.renderer + return None + async def do_log_stats(self) -> None: pass diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index 287f12b9e..b557c07dd 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -116,7 +116,6 @@ class AsyncOmniLLM(AsyncLLM): # InputProcessor (converts Inputs --> EngineCoreRequests). self.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=tokenizer, mm_registry=mm_registry, ) diff --git a/vllm_omni/entrypoints/chat_utils.py b/vllm_omni/entrypoints/chat_utils.py index 9a8313dba..ca15305c5 100644 --- a/vllm_omni/entrypoints/chat_utils.py +++ b/vllm_omni/entrypoints/chat_utils.py @@ -39,17 +39,27 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser): self._mm_processor_kwargs = mm_processor_kwargs def parse_video(self, video_url: str | None, uuid: str | None = None) -> None: - video = self._connector.fetch_video_async(video_url=video_url) if video_url else None - - placeholder = self._tracker.add("video", video, uuid) + # OMNI: Follow upstream async pattern - create coroutine that resolves to (data, uuid) + coro = self._video_with_uuid_async(video_url, uuid) + placeholder = self._tracker.add("video", coro) self._add_placeholder("video", placeholder) # Extract audio from video if use_audio_in_video is True if video_url and self._mm_processor_kwargs and self._mm_processor_kwargs.get("use_audio_in_video", False): - audio_coro = self._extract_audio_from_video_async(video_url) - audio_placeholder = self._tracker.add("audio", audio_coro, uuid) + audio_coro = self._audio_from_video_with_uuid_async(video_url, uuid) + audio_placeholder = self._tracker.add("audio", audio_coro) self._add_placeholder("audio", audio_placeholder) + async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None): + """Fetch video and return (video, uuid) tuple.""" + video = await self._connector.fetch_video_async(video_url=video_url) if video_url else None + return video, uuid + + async def _audio_from_video_with_uuid_async(self, video_url: str, uuid: str | None): + """Extract audio from video and return (audio, uuid) tuple.""" + audio = await self._extract_audio_from_video_async(video_url) + return audio, uuid + async def _extract_audio_from_video_async(self, video_url: str) -> tuple[np.ndarray, int | float]: """ Extract audio from video URL using librosa. @@ -132,9 +142,17 @@ def parse_chat_messages_futures( mm_processor_kwargs: dict[str, Any] | None = None, ) -> tuple[ list[ConversationMessage], - Awaitable[MultiModalDataDict | None], - MultiModalUUIDDict | None, + Awaitable[tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]], ]: + """Parse chat messages and return conversation with multimodal data future. + + OMNI: Updated to use upstream vLLM v0.15.0 API where resolve_items() + returns both mm_data and mm_uuids together as a tuple. + + Returns: + Tuple of (conversation, mm_future) where mm_future resolves to + (mm_data, mm_uuids) when awaited. + """ conversation: list[ConversationMessage] = [] mm_tracker = OmniAsyncMultiModalItemTracker(model_config) @@ -155,7 +173,8 @@ def parse_chat_messages_futures( _postprocess_messages(conversation) - return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids() + # OMNI: Use upstream resolve_items() which returns (mm_data, mm_uuids) tuple + return conversation, mm_tracker.resolve_items() def _parse_chat_message_content( diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 6bb8ac366..0356ea231 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -161,9 +161,7 @@ class OmniLLM(LLM): log_stats=self.llm_engine.log_stats, engine_core_output_type=engine_args.engine_output_type, ) - self.llm_engine.input_processor = OmniInputProcessor( - vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer - ) + self.llm_engine.input_processor = OmniInputProcessor(vllm_config=self.llm_engine.vllm_config) self.engine_class = type(self.llm_engine) self.request_counter = Counter() diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index fc01c147c..133e9ec84 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -670,12 +670,22 @@ def _stage_worker( break lock_files = acquired_lock_fds + + # Set FD_CLOEXEC on all lock file descriptors to prevent child processes + # (e.g., EngineCore) from inheriting them, which would cause deadlock + for lock_fd in acquired_lock_fds: + try: + flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) + fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + except (OSError, ValueError): + pass except Exception as e: logger.debug( "[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e, ) + # Init engine based on stage_type logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) if engine_args.get("async_chunk", False): @@ -686,30 +696,29 @@ def _stage_worker( break engine_args["stage_connector_spec"] = stage_connector_spec engine_args["stage_id"] = stage_id - try: - if stage_type == "diffusion": - 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) - finally: - # Release all locks by closing file descriptors - # Locks are automatically released when file descriptors are closed - # or when process dies - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass + if stage_type == "diffusion": + 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) + + # Release all locks AFTER engine initialization completes + for lock_fd in lock_files: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + _os.close(lock_fd) + logger.debug("Released initialization lock (fd=%s)", lock_fd) + except (OSError, ValueError): + pass + lock_files = [] # Clear after release + logger.debug("Engine initialized") # Initialize OmniConnectors if configured connectors: dict[tuple[str, str], OmniConnectorBase] | None = {} diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 9620f3bef..b81f8a1f8 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -362,7 +362,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer, model_config=model_config, ) - conversation, mm_data_future, mm_uuids = parse_chat_messages_futures( + # OMNI: Updated for vLLM v0.15.0 API - resolve_items() returns (mm_data, mm_uuids) tuple + conversation, mm_future = parse_chat_messages_futures( messages, model_config, content_format=resolved_content_format, @@ -409,7 +410,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): **hf_chat_template_kwargs, ) - mm_data = await mm_data_future + # OMNI: Await the combined future to get both mm_data and mm_uuids + mm_data, mm_uuids = await mm_future # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser @@ -1677,7 +1679,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output - audio_data = final_res.multimodal_output.get("audio") + # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput + # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421 + audio_data = final_res.outputs[0].multimodal_output.get("audio") if stream: audio_tensor = audio_data[-1].float().detach().cpu().numpy() else: @@ -1763,9 +1767,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if omni_outputs.images: images = omni_outputs.images # Fall back to request_output for pipeline mode - elif final_res is not None: - if hasattr(final_res, "multimodal_output") and final_res.multimodal_output: - image_data = final_res.multimodal_output.get("image") + # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput + elif final_res is not None and final_res.outputs: + completion_output = final_res.outputs[0] + if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output: + image_data = completion_output.multimodal_output.get("image") if image_data is not None: if isinstance(image_data, Image.Image): images.append(image_data) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 673cf6e47..a4d0432c5 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -105,6 +105,9 @@ class OmniEmbedsPrompt(EmbedsPrompt): # Must ensure that all additional prompt types are inherited from vLLM prompt types # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime. # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate() +OmniSingletonPrompt: TypeAlias = str | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt +"""Omni singleton prompt type extending vLLM's SingletonPrompt with additional fields.""" + OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py index f21af5d2c..694fd44bf 100644 --- a/vllm_omni/inputs/parse.py +++ b/vllm_omni/inputs/parse.py @@ -4,13 +4,12 @@ from vllm.inputs.parse import ( ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt, + SingletonPrompt, ) -from vllm_omni.inputs.data import OmniSingletonPrompt - def parse_singleton_prompt_omni( - prompt: OmniSingletonPrompt, + prompt: SingletonPrompt, ) -> ParsedSingletonPrompt: """Parse a singleton prompt into a typed parsed prompt. diff --git a/vllm_omni/model_executor/layers/rotary_embedding/__init__.py b/vllm_omni/model_executor/layers/rotary_embedding/__init__.py new file mode 100644 index 000000000..02ebeb194 --- /dev/null +++ b/vllm_omni/model_executor/layers/rotary_embedding/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .mrope import OmniMRotaryEmbedding + +__all__ = ["OmniMRotaryEmbedding"] diff --git a/vllm_omni/model_executor/layers/rotary_embedding/mrope.py b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py new file mode 100644 index 000000000..463e55507 --- /dev/null +++ b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py @@ -0,0 +1,554 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Omni-extended MRotaryEmbedding with multimodal position computation methods. + +This module extends the upstream vLLM MRotaryEmbedding with additional methods +for computing input positions for various multimodal scenarios including: +- Image/Video inputs (Qwen2.5-VL style) +- Audio inputs (Qwen2.5-Omni style) +- Audio-in-video interleaved inputs +- GLM4V style inputs +""" + +import itertools + +import torch +from transformers import PretrainedConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding + +logger = init_logger(__name__) + + +class OmniMRotaryEmbedding(MRotaryEmbedding): + """Omni-extended MRotaryEmbedding with multimodal position computation. + + Extends the upstream MRotaryEmbedding with additional class methods for + computing input positions for various multimodal scenarios. + """ + + @classmethod + def get_input_positions( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor | None, + video_grid_thw: list[list[int]] | torch.Tensor | None, + second_per_grid_ts: list[float] | None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[list[list[int]], int]: + """Get mrope input positions and delta value.""" + + image_grid_thw = [] if image_grid_thw is None else image_grid_thw + video_grid_thw = [] if video_grid_thw is None else video_grid_thw + second_per_grid_ts = [] if second_per_grid_ts is None else second_per_grid_ts + + llm_positions, mrope_position_delta = cls.get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + return llm_positions.tolist(), mrope_position_delta + + @classmethod + def get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float], + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + from vllm.transformers_utils.config import thinker_uses_mrope + + if thinker_uses_mrope(hf_config): + return cls._omni_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + elif hf_config.model_type in ["glm4v", "glm4v_moe"]: + return cls._glm4v_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + context_len=context_len, + seq_len=seq_len, + ) + else: + return cls._vl_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + ) + + @classmethod + def _glm4v_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + context_len: int = 0, + seq_len: int | None = None, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value for GLM4V.""" + + image_token_id = hf_config.image_token_id + video_start_token_id = hf_config.video_start_token_id + video_end_token_id = hf_config.video_end_token_id + spatial_merge_size = hf_config.vision_config.spatial_merge_size + llm_pos_ids_list: list = [] + + if not (image_grid_thw is None and video_grid_thw is None): + if isinstance(image_grid_thw, torch.Tensor): + image_grid_thw = image_grid_thw.tolist() + + input_token_type: list[str] = [] + video_check_flg = False + for token in input_tokens: + if token == video_start_token_id: + video_check_flg = True + elif token == video_end_token_id: + video_check_flg = False + + if (token == image_token_id) and (video_check_flg is False): + input_token_type.append("image") + elif (token == image_token_id) and (video_check_flg is True): + input_token_type.append("video") + else: + input_token_type.append("text") + + input_type_group: list[tuple[str, int, int]] = [] + for key, group_iter in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + group_list = list(group_iter) + start_index = group_list[0][0] + end_index = group_list[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + video_frame_num = 1 + mm_data_idx = 0 + for modality_type, start_idx, end_idx in input_type_group: + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + if modality_type == "image": + t, h, w = ( + image_grid_thw[mm_data_idx][0], + image_grid_thw[mm_data_idx][1], + image_grid_thw[mm_data_idx][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = ( + t, + h // spatial_merge_size, + w // spatial_merge_size, + ) + + t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) + mm_data_idx += 1 + + elif modality_type == "video": + t, h, w = ( + video_frame_num, + image_grid_thw[mm_data_idx][1], + image_grid_thw[mm_data_idx][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = ( + t, + h // spatial_merge_size, + w // spatial_merge_size, + ) + + for t_idx in range(llm_grid_t): + t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) + + mm_data_idx += 1 + video_frame_num += 1 + + else: + text_len = end_idx - start_idx + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + video_frame_num = 1 + + else: + text_len = len(input_tokens) + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1)) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + llm_positions = llm_positions[:, context_len:seq_len] + mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() + return llm_positions, mrope_position_delta + + @classmethod + def _vl_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float], + context_len: int = 0, + seq_len: int | None = None, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value.""" + + image_token_id = hf_config.image_token_id + video_token_id = hf_config.video_token_id + vision_start_token_id = hf_config.vision_start_token_id + spatial_merge_size = hf_config.vision_config.spatial_merge_size + tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0) + + input_tokens_tensor = torch.tensor(input_tokens) + vision_start_indices = torch.argwhere(input_tokens_tensor == vision_start_token_id).squeeze(1) + vision_tokens = input_tokens_tensor[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + llm_pos_ids_list: list = [] + + st = 0 + remain_images, remain_videos = image_nums, video_nums + + image_index, video_index = 0, 0 + for _ in range(image_nums + video_nums): + video_second_per_grid_t = 0.0 + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_second_per_grid_t = 1.0 + if second_per_grid_ts: + video_second_per_grid_t = second_per_grid_ts[video_index] + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = ( + t, + h // spatial_merge_size, + w // spatial_merge_size, + ) + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + t_index = ( + ( + torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w) + * video_second_per_grid_t + * tokens_per_second + ) + .long() + .flatten() + ) + + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() + llm_positions = llm_positions[:, context_len:seq_len] + + return llm_positions, mrope_position_delta + + @classmethod + def _omni_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value (Qwen2.5-Omni version). + + Differences from MRotaryEmbedding: + 1. Add audio support (and related `audio_feature_lengths`). + 2. Add `use_audio_in_video` option to read audio from video inputs. + In this case, audio and vision position ids will be split into + chunks and interleaved. + + Example: + + (V_i are vision position ids, A_i are audio position ids) + + |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|... + |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |... + """ + + thinker_config = hf_config.thinker_config + try: + audio_token_id = thinker_config.audio_token_index + image_token_id = thinker_config.image_token_index + video_token_id = thinker_config.video_token_index + except Exception: + logger.info("Multimodal token idx changed!") + audio_token_id = thinker_config.audio_token_id + image_token_id = thinker_config.image_token_id + video_token_id = thinker_config.video_token_id + + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + vision_start_token_id = thinker_config.vision_start_token_id + vision_end_token_id = thinker_config.vision_end_token_id + seconds_per_chunk = thinker_config.seconds_per_chunk + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25) + + if isinstance(image_grid_thw, list): + image_grid_thw = torch.tensor(image_grid_thw) + if isinstance(video_grid_thw, list): + video_grid_thw = torch.tensor(video_grid_thw) + + src_item = input_tokens + audio_seqlens = audio_feature_lengths + if not second_per_grid_ts: + second_per_grid_ts = [1] * video_grid_thw.shape[0] + audio_idx = 0 + video_idx = 0 + image_idx = 0 + new_src_item: list[int] = [] + llm_pos_ids_list: list[torch.Tensor] = [] + + idx = 0 + while idx < len(src_item): + new_src_item_len = len(new_src_item) + start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]: + if use_audio_in_video and idx > 0: + if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id: + # processing the <|audio_eos|> before <|vision_eos|> + start_idx -= 1 + elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id: + # processing the <|audio_bos|> after <|vision_eos|> + start_idx -= 1 + new_src_item.append(src_item[idx]) + llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1) + llm_pos_ids_list.append(llm_pos_ids) + elif src_item[idx] == audio_token_id: + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1 + new_src_item.extend([audio_token_id] * place_num) + llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx + llm_pos_ids_list.append(llm_pos_ids) + audio_idx += 1 + elif src_item[idx] == image_token_id: + grid_t = image_grid_thw[image_idx][0] + grid_hs = image_grid_thw[:, 1] + grid_ws = image_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long() + llm_pos_ids = cls._get_llm_pos_ids_for_vision( + start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2) + new_src_item.extend([image_token_id] * vision_seqlen) + image_idx += 1 + elif src_item[idx] == video_token_id and not use_audio_in_video: + grid_t = video_grid_thw[video_idx][0] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() + llm_pos_ids = cls._get_llm_pos_ids_for_vision( + start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + new_src_item.extend([video_token_id] * vision_seqlen) + video_idx += 1 + else: + # read audio from video + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + grid_t = video_grid_thw[video_idx][0] + grid_h = video_grid_thw[video_idx][1] + grid_w = video_grid_thw[video_idx][2] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) + t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() + t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk) + place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2 + pure_audio_len = place_num - 2 + added_audio_len = 0 + audio_llm_pos_ids_list: list[torch.Tensor] = [] + for t_chunk in t_index_split_chunk: + vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2) + new_src_item.extend([video_token_id] * vision_ntoken_per_chunk) + vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision( + start_idx, + video_idx, + spatial_merge_size, + t_chunk, + grid_hs, + grid_ws, + ).split(1, dim=1) + llm_pos_ids_list.extend(vision_llm_pos_ids_list) + new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id]) + audio_start_idx = ( + start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1 + ) + if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0: + audio_llm_pos_ids_list = ( + torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1) + + audio_start_idx + ).split(1, dim=1) + else: + audio_llm_pos_ids_list = [] + added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + if added_audio_len < pure_audio_len: + new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id]) + audio_llm_pos_ids_list = ( + torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1 + ).split(1, dim=1) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + audio_idx += 1 + video_idx += 1 + # move to the next token + idx += len(new_src_item) - new_src_item_len + + llm_positions = torch.cat(llm_pos_ids_list, dim=1) + mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item) + llm_positions = llm_positions[:, context_len:seq_len] + + return llm_positions, mrope_position_delta + + @staticmethod + def _get_llm_pos_ids_for_vision( + start_idx: int, + vision_idx: int, + spatial_merge_size: int, + t_index: list[int], + grid_hs: torch.Tensor, + grid_ws: torch.Tensor, + ) -> torch.Tensor: + llm_pos_ids_list = [] + llm_grid_h = grid_hs[vision_idx] // spatial_merge_size + llm_grid_w = grid_ws[vision_idx] // spatial_merge_size + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten() + t_index_tensor = ( + torch.Tensor(t_index).to(llm_grid_h.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten() + ) + _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index]) + llm_pos_ids_list.append(_llm_pos_ids + start_idx) + llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1) + return llm_pos_ids + + @staticmethod + def _split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]: + ranges: list[list[int]] = [[] for _ in range((max(lst) // interval) + 1)] + for num in lst: + index = num // interval + ranges[index].append(num) + return ranges + + @classmethod + def omni_get_updates_use_audio_in_video( + cls, + thinker_config: PretrainedConfig, + audio_len: int, + video_grid_thw: list[int] | torch.Tensor, + video_second_per_grid_t: float, + ) -> list[int]: + """Get video prompt updates when `use_audio_in_video` is True. + + In this case, audio and vision update ids will be split into + chunks and interleaved (details in `_omni_get_input_positions_tensor`). + + <|video_bos|><|VIDEO|><|video_eos|> => + <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|> + """ + + audio_token_id = thinker_config.audio_token_index + video_token_id = thinker_config.video_token_index + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + seconds_per_chunk = thinker_config.seconds_per_chunk + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25) + + grid_t = video_grid_thw[0] + grid_h = video_grid_thw[1] + grid_w = video_grid_thw[2] + t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) + t_index = (torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second).long() + t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk) + + updates = [audio_start_token_id] + added_audio_len = 0 + for t_chunk in t_index_split_chunk: + vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2) + updates.extend([video_token_id] * vision_ntoken_per_chunk) + + audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len) + updates.extend(audio_chunk_size * [audio_token_id]) + added_audio_len += audio_chunk_size + if added_audio_len < audio_len: + updates.extend((audio_len - added_audio_len) * [audio_token_id]) + updates.extend([audio_end_token_id]) + + return updates diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index de706ea9a..24f2aecc1 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -609,7 +609,8 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - audio_features = audio_outputs.last_hidden_state + # OMNI: audio_tower.forward() returns hidden_states tensor directly + audio_features = audio_outputs return audio_features.split(audio_output_lengths.tolist()) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 9fb4249e6..90c718d40 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,10 +1,9 @@ import sys from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt - -# from vllm.model_executor.layers.rotary_embedding import ( -# MRotaryEmbedding as _OriginalMRotaryEmbedding, -# ) +from vllm.model_executor.layers.rotary_embedding import ( + MRotaryEmbedding as _OriginalMRotaryEmbedding, +) from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest @@ -13,6 +12,7 @@ from vllm.v1.request import Request as _OriginalRequest import vllm_omni.logger # noqa: F401 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest for module_name, module in sys.modules.items(): @@ -25,8 +25,8 @@ for module_name, module in sys.modules.items(): module.EngineCoreOutputs = OmniEngineCoreOutputs if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt: module.TokensPrompt = OmniTokensPrompt - # if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding: - # module.MRotaryEmbedding = MRotaryEmbedding + if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding: + module.MRotaryEmbedding = OmniMRotaryEmbedding if hasattr(module, "Request") and module.Request == _OriginalRequest: module.Request = OmniRequest if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 1c64cd69a..9332abeb1 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -219,6 +219,24 @@ class GPUARModelRunner(OmniGPUModelRunner): use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + # True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec + + has_separate_kv_update = not all( + all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id]) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded, + num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, @@ -230,6 +248,7 @@ class GPUARModelRunner(OmniGPUModelRunner): use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, + slot_mappings=slot_mappings_by_group, ) ( @@ -260,6 +279,7 @@ class GPUARModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index b6f122a38..9125cc52a 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -189,6 +189,24 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + # OMNI: True if any attention backend handles KV cache update separately + # from forward() (i.e., forward_includes_kv_cache_update=False). When true, + # slot_mappings must use padded dimensions to match the key/value tensors. + from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec + + has_separate_kv_update = not all( + all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id]) + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups) + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec) + ) + + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded, + num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs), + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, @@ -200,6 +218,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, + slot_mappings=slot_mappings_by_group, ) ( @@ -234,6 +253,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, @@ -539,6 +559,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): attn_metadata: PerLayerAttnMetadata | None = None + # OMNI: Get slot mappings before building attention metadata + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: @@ -564,6 +592,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): max_query_len=max_query_len, ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, ) with self.maybe_dummy_run_with_lora( @@ -628,6 +657,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations ), ): outputs = self.model( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e3a3769d8..14080cf14 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -552,6 +552,13 @@ class OmniGPUModelRunner(GPUModelRunner): attn_metadata: PerLayerAttnMetadata | None = None + slot_mappings_by_group, slot_mappings = self._get_slot_mappings( + num_tokens_padded=num_tokens, + num_reqs_padded=num_reqs_padded, + num_tokens_unpadded=num_tokens_unpadded, + ubatch_slices=ubatch_slices_padded, + ) + # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: @@ -577,6 +584,7 @@ class OmniGPUModelRunner(GPUModelRunner): max_query_len=max_query_len, ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, + slot_mappings=slot_mappings_by_group, ) with self.maybe_dummy_run_with_lora( @@ -641,6 +649,7 @@ class OmniGPUModelRunner(GPUModelRunner): cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, ), ): if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): @@ -688,6 +697,7 @@ class OmniGPUModelRunner(GPUModelRunner): num_tokens, use_cudagraphs=use_cudagraphs, is_graph_capturing=is_graph_capturing, + slot_mappings=slot_mappings, ) # We register layerwise NVTX hooks here after the first dynamo tracing is