diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 54ed64261..0c06e3adc 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -32,6 +32,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - By default, image tasks talk to `http://:/v1/chat/completions`; video tasks talk to `/v1/videos`. - If you run the server on another host or port, pass `--base-url` accordingly. +- To compare diffusion multi-replica serving, restart the server with different + `runtime.num_replicas` values and keep the benchmark parameters unchanged. ## 2. Supported Datasets @@ -93,6 +95,35 @@ You can point to your own trace using `--dataset-path`. ## 3. Benchmark Parameters +### Multi-replica sweep + +For replica scaling, use the same model, prompt count, image/video size, and +step count for every run. Change only the serving stage config: + +```yaml +runtime: + devices: "0,1,2,3" + num_replicas: 4 +``` + +Then run a fixed benchmark and write one JSON file per replica count: + +```bash +python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8099 \ + --model Qwen/Qwen-Image \ + --task t2i \ + --dataset random \ + --num-prompts 32 \ + --max-concurrency 8 \ + --request-rate inf \ + --num-inference-steps 10 \ + --width 512 --height 512 \ + --output-file diffusion_replica_4.json +``` + +Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`. + ### Basic flags - `--base-url`: Server address; `--endpoint` selects the path appended to this base URL. diff --git a/docs/design/feature/qwen3_omni_thinker_mr_v2.md b/docs/design/feature/qwen3_omni_thinker_mr_v2.md new file mode 100644 index 000000000..758fd2051 --- /dev/null +++ b/docs/design/feature/qwen3_omni_thinker_mr_v2.md @@ -0,0 +1,1153 @@ +# 设计文档：Qwen3-Omni Thinker 迁移至 Model Runner V2 + +**作者：**（基于 RFC #1770 决策自动生成） +**状态：** 草案 +**相关文档：** [RFC #1770 — 迁移至 Model Runner V2](https://github.com/vllm-project/vllm-omni/issues/1770)、[RFC #1595 — 解耦模型特定逻辑](https://github.com/vllm-project/vllm-omni/issues/1595) +**范围：** Qwen3-Omni-30B-A3B-Instruct **仅 Thinker 阶段** 迁移至 MR V2 +**目标时间线：** 2026 年 5 月设为默认 + +--- + +## 1. 目标 + +将 Qwen3-Omni Thinker（AR 文本生成阶段）从 v1 的 `OmniGPUModelRunner → GPUARModelRunner` 继承链迁移至 v2 的 `GPUModelRunner` + `ModelState` 架构。Thinker 是首个落地 MR V2 的模型，将作为后续 Talker、TTS、Code2Wav 等模型迁移的模板。 + +### 非目标（本文档不涉及） + +- Talker / TTS / Code2Wav / Diffusion runner 迁移（后续阶段）。 +- 修改上游 `GPUModelRunner` 或 `ModelState` 接口。 +- NPU / XPU 平台适配。 + +### 约束条件 + +- **不修改上游 v2 代码。** 所有扩展通过继承、组合和 `init_model_state()` 注册机制实现。 +- Talker、TTS、Code2Wav 必须在同一抽象层上保持可行性。设计决策不能为后续阶段制造障碍。 +- 过渡期间 v1 和 v2 runner 通过 `VLLM_OMNI_USE_V2_RUNNER` 环境变量共存。 + +--- + +## 2. 背景：当前 v1 架构（Thinker 路径） + +### 2.1 类继承层次 + +``` +vllm.v1.worker.gpu_model_runner.GPUModelRunner (v1, ~6283 行) + └── OmniGPUModelRunner (~1403 行) + └── GPUARModelRunner (~677 行) +``` + +### 2.2 Thinker 阶段实际需要的功能 + +Qwen3-Omni Thinker 是一个多模态编码器-解码器（视觉/音频 Transformer + MoE LLM）。在 v1 runner 中，Thinker 路径仅使用了 `OmniGPUModelRunner` 功能的**子集**： + +| 功能 | Thinker 是否使用？ | 说明 | +|------|:-:|------| +| `model_intermediate_buffer` | 部分 | 在请求进入时填充（prompt_embeds、additional_information）；每步以 `runtime_additional_information` 形式聚合——但 **Thinker 的 `forward()` 会忽略它**（被 `**kwargs` 吸收）。该缓冲区的存在是为了下游的 Talker/Code2Wav 阶段。 | +| `prompt_embeds` 覆盖 | 否 | Thinker 的 `has_preprocess = False`；`prompt_embeds_cpu` 会被解码并存储，但不会进行覆盖。 | +| M-RoPE 位置编码 | **是** | Thinker 实现了 `SupportsMRoPE` 接口；`get_mrope_input_positions()` 从 `image_grid_thw`、`video_grid_thw`、`audio_feature_lengths` 计算 3D 位置。 | +| `OmniOutput` 包装 | **是** | `make_omni_output()` 将 Thinker 输出 `(text_hidden_states, captured_layer_dict)` 包装为 `OmniOutput`，附带多模态侧载信息（TTS embed BOS/EOS/PAD）。 | +| `extract_multimodal_outputs` | **是** | 将 `OmniOutput` 拆分为 `text_hidden_states` + `multimodal_outputs` 字典。 | +| `has_preprocess` / 自定义预处理 | 否 | Thinker：`has_preprocess = False`。 | +| `has_postprocess` / 自定义后处理 | 否 | Thinker：`has_postprocess = False`。 | +| `talker_mtp` 子模块 | 否 | Thinker：`self.talker = None`，无 talker_mtp。 | +| `pooler_output`（hidden + mm） | **是** | AR runner 将每个请求的 hidden states + 多模态输出切片后组装为 `pooler_output`，供下游阶段使用。 | +| KV 传输 | **是** | `OmniKVTransferManager` 为已完成请求提取 KV 缓存。 | +| Deepstack 视觉 | **是** | Thinker 模型 `forward()` 内部处理，对 runner 透明。 | + +**关键洞察：** Thinker 是最简单的 Omni 模型阶段。它不使用 `preprocess`/`postprocess`、`talker_mtp` 或 `prompt_embeds` 覆盖。它**需要**的是 M-RoPE、`OmniOutput` 处理、`pooler_output` 构建和 KV 传输。 + +--- + +## 3. 设计概览 + +### 3.1 目标架构 + +``` +vllm.v1.worker.gpu.model_runner.GPUModelRunner (v2, ~1168 行, 不修改) + └── OmniGPUModelRunner (薄继承层, ~150 行) + └── OmniARModelRunner (AR 专用, ~350 行) + +vllm.v1.worker.gpu.model_states.interface.ModelState (v2 接口, 不修改) + └── DefaultModelState (v2 默认实现) + └── OmniModelState (通用 Omni 基类, ~300 行) + +init_model_state() ← 通过架构名检查扩展 (v2 工厂函数, 新增 ~5 行) + +OmniIntermediateBuffer (新增, ~120 行) +OmniModelStatePlugin (新增, ABC, ~30 行) +PositionEncodingStrategy (新增, ABC + 4 种实现, ~200 行) +OmniKVConnector (新增, 封装 OmniKVTransferManager, ~80 行) +``` + +### 3.2 设计决策（来自 RFC #1770 开放问题） + +| # | 决策 | 理由 | +|---|------|------| +| 1 | 通过 `init_model_state()` 工厂函数注册 `OmniModelState` | 比 runner 侧注入更清晰；`init_model_state()` 已有基于架构名的分发逻辑（Whisper），将 Omni 架构加入同一模式。 | +| 2 | 对大型中间张量使用 `UvaBackedTensor` | 减少 `prompt_embeds` 和捕获的 hidden states 的显式 CPU↔GPU 拷贝。与 v2 的内存管理方式一致。 | +| 3 | Generation runner 复用 v2 的 `pool()` 代码路径 | Generation 阶段无需采样；`pool()` 已绕过 `sample_tokens`。适配 `pool()` 的返回语义以匹配 Generation 输出。 | +| 4 | V2 runner 于 2026 年 5 月成为默认 | 功能开关 `VLLM_OMNI_USE_V2_RUNNER` 翻转默认值。V1 保留一个发布周期后移除。 | + +--- + +## 4. 详细设计 + +### 4.1 `init_model_state()` 扩展 + +v2 工厂函数位于 `vllm/v1/worker/gpu/model_states/__init__.py`，当前按架构名分发。我们在 `vllm_omni` 中添加 Omni 感知的包装函数来**扩展**（而非修改）它： + +```python +# vllm_omni/worker_v2/model_states/__init__.py + +from vllm.v1.worker.gpu.model_states import init_model_state as _upstream_init + +# 需要使用 OmniModelState 的架构列表 +_OMNI_ARCHITECTURES = { + "Qwen3OmniMoeForConditionalGeneration", + "MammothModa2ForConditionalGeneration", + "MiMoAudioForConditionalGeneration", + # 后续模型在此添加 +} + +def init_omni_model_state(vllm_config, model, encoder_cache, device): + """扩展 init_model_state，增加 Omni 架构分发。""" + archs = set(vllm_config.model_config.architectures or []) + if archs & _OMNI_ARCHITECTURES: + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + return OmniModelState(vllm_config, model, encoder_cache, device) + return _upstream_init(vllm_config, model, encoder_cache, device) +``` + +`OmniGPUModelRunner` 通过覆写 `load_model` 使用 `init_omni_model_state` 替代上游工厂函数。这**不需要任何上游代码修改**。 + +### 4.2 `OmniModelState` — 通用基类 + +`OmniModelState` 继承 `DefaultModelState`。它与模型无关：不包含 Qwen3 特定代码，也不包含 Talker 逻辑。它处理**所有** Omni 模型共享的横切关注点。 + +```python +class OmniModelState(DefaultModelState): + """通用 Omni ModelState — 适用于所有 Omni 模型阶段。 + + 职责： + - 跨阶段中间缓冲区（prompt_embeds、additional_information） + - OmniOutput → text_hidden + multimodal_outputs 提取 + - model_intermediate_buffer / runtime_additional_information 注入 + - 插件生命周期分发 + """ + + def __init__(self, vllm_config, model, encoder_cache, device): + super().__init__(vllm_config, model, encoder_cache, device) + max_num_reqs = self.scheduler_config.max_num_seqs + self.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) + self.has_preprocess = getattr(model, "has_preprocess", False) + self.has_postprocess = getattr(model, "has_postprocess", False) + self.have_multimodal_outputs = getattr(model, "have_multimodal_outputs", False) + self.plugins: list[OmniModelStatePlugin] = [] + + # 允许模型在初始化时注册插件 + if hasattr(model, "get_omni_plugins"): + for plugin in model.get_omni_plugins(): + self.register_plugin(plugin) + + def register_plugin(self, plugin: "OmniModelStatePlugin") -> None: + self.plugins.append(plugin) + + # --- 请求生命周期 --- + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + self.intermediate_buffer.add_request(req_index, new_req_data) + for plugin in self.plugins: + plugin.on_add_request(req_index, new_req_data) + + def remove_request(self, req_index: int) -> None: + self.intermediate_buffer.remove_request(req_index) + for plugin in self.plugins: + plugin.on_remove_request(req_index) + + # --- 输入准备 --- + + def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]: + base = super().prepare_inputs(input_batch, req_states) + + # 将中间缓冲区作为模型 kwargs 注入 + # v2 通过 **展开 将此字典合并到 model_inputs 中 + buffer_list = self.intermediate_buffer.gather(input_batch) + base["model_intermediate_buffer"] = buffer_list + base["runtime_additional_information"] = buffer_list # 向后兼容 + + # 插件提供的额外输入 + for plugin in self.plugins: + base.update(plugin.prepare_extra_inputs(input_batch, req_states)) + + return base + + # --- 输出处理 --- + + def postprocess_model_output(self, model_output, input_batch, req_states): + """将原始模型输出转换为 (text_hidden, multimodal_outputs)。 + + 处理 OmniOutput 解包和 make_omni_output 转换。 + 由 OmniARModelRunner.sample_tokens() 调用。 + """ + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + buffer_list = self.intermediate_buffer.gather(input_batch) + model_output = self.model.make_omni_output( + model_output, + model_intermediate_buffer=buffer_list, + runtime_additional_information=buffer_list, + ) + + if self.have_multimodal_outputs and isinstance(model_output, OmniOutput): + text_hidden = model_output.text_hidden_states + multimodal_outputs = model_output.multimodal_outputs or {} + elif isinstance(model_output, (list, tuple)): + text_hidden = model_output[0] + multimodal_outputs = {} + else: + text_hidden = model_output + multimodal_outputs = {} + + # 运行插件（例如 Talker 阶段的 TalkerMTPPlugin —— Thinker 不涉及） + for plugin in self.plugins: + text_hidden, multimodal_outputs = plugin.postprocess( + text_hidden, multimodal_outputs, input_batch, req_states + ) + + return text_hidden, multimodal_outputs + + def prepare_dummy_inputs(self, num_reqs, num_tokens) -> dict[str, Any]: + base = super().prepare_dummy_inputs(num_reqs, num_tokens) + base["model_intermediate_buffer"] = [{} for _ in range(num_reqs)] + base["runtime_additional_information"] = base["model_intermediate_buffer"] + return base +``` + +#### 4.2.1 对 Thinker 的影响 + +对于 Thinker 阶段（`has_preprocess=False`、`has_postprocess=False`、无插件）： + +- `add_request` → 将 `prompt_embeds` 和 `additional_information` 存入 `intermediate_buffer`（供下游使用）；通过 `super().add_request()` 调用 `rope_state.init_prefill_positions()` 初始化 M-RoPE。 +- `prepare_inputs` → 返回 `{"positions": mrope_positions, "model_intermediate_buffer": [...], "runtime_additional_information": [...]}`。Thinker 的 `forward()` 通过 `**kwargs` 静默吸收这些参数。 +- `postprocess_model_output` → 通过 `make_omni_output()` 将 `(text_hidden_states, captured_layer_dict)` 转换为 `OmniOutput`，然后提取 `text_hidden_states` + `multimodal_outputs`（捕获的层、TTS embeds）。 +- 无插件执行。 + +### 4.3 `OmniIntermediateBuffer` + +基于索引设计，与 v2 的 `RequestState` 槽位管理对齐。对 `prompt_embeds` 使用 `UvaBackedTensor`。 + +```python +class OmniIntermediateBuffer: + """多阶段流水线的每请求中间状态。 + + 使用 req_index（而非 req_id）进行 O(1) 访问，与 v2 RequestState 对齐。 + """ + + def __init__(self, max_num_reqs: int): + self.buffers: list[dict[str, Any]] = [{} for _ in range(max_num_reqs)] + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + """解码并存储 prompt_embeds 和 additional_information。""" + info: dict[str, Any] = {} + + # 解码 prompt_embeds（序列化张量载荷） + pe = getattr(new_req_data, "prompt_embeds", None) + if pe is not None: + pe_cpu = _resolve_prompt_embeds(pe) + if pe_cpu is not None: + info["prompt_embeds_cpu"] = pe_cpu + + # 解码 additional_information（字典或旧版载荷） + ai = getattr(new_req_data, "additional_information", None) + if ai is not None: + info.update(_resolve_additional_information(ai)) + + # 存储 mm_features 引用，供需要的模型使用（如 MiMoAudio） + if new_req_data.mm_features: + info["mm_features"] = new_req_data.mm_features + + # 存储 req_id，供需要的模型使用 + info["req_id"] = new_req_data.req_id + + self.buffers[req_index] = info + + def remove_request(self, req_index: int) -> None: + self.buffers[req_index] = {} + + def gather(self, input_batch: InputBatch) -> list[dict]: + """按批次顺序聚合中间信息，使用 idx_mapping。""" + return [self.buffers[idx] for idx in input_batch.idx_mapping_np] + + def update(self, req_index: int, updates: dict) -> None: + """合并更新。张量自动 detach；GPU 常驻键保留在 GPU 上。""" + existing = self.buffers[req_index] + for k, v in updates.items(): + if isinstance(v, torch.Tensor): + existing[k] = v.detach().clone() + else: + existing[k] = v +``` + +#### 4.3.1 Thinker 的缓冲区生命周期 + +``` + 请求到达 + │ + ▼ + add_request(req_index, new_req_data) + ├── 解码 prompt_embeds → info["prompt_embeds_cpu"] + ├── 解码 additional_information → info[...] + ├── info["mm_features"] = new_req_data.mm_features + └── info["req_id"] = new_req_data.req_id + │ + ▼ （每步执行） + gather(input_batch) → list[dict] → model_intermediate_buffer kwarg + │ （Thinker 通过 **kwargs 忽略） + ▼ + 请求完成 + │ + ▼ + remove_request(req_index) → buffers[i] = {} +``` + +对于 Thinker，缓冲区虽被填充但**模型不消费**。它的存在是为了下游的 Talker/Code2Wav 阶段能够通过跨阶段 KV 传输元数据接收数据。 + +### 4.4 `OmniModelStatePlugin` — 扩展点 + +模型特定逻辑的抽象接口。**Thinker 不注册任何插件。** 此接口为 Talker（`TalkerMTPPlugin`）和未来模型而设。 + +```python +class OmniModelStatePlugin(ABC): + """OmniModelState 中模型特定行为的扩展点。 + + 生命周期： + on_add_request() — 新请求进入批次 + on_remove_request() — 请求完成 + prepare_extra_inputs() — 向 model_inputs 注入额外 kwargs + postprocess() — 处理模型输出（如 talker_mtp 前向传播） + dummy_run() — 内存分析 / CUDA 图捕获 + """ + + @abstractmethod + def on_add_request(self, req_index, new_req_data) -> None: ... + + @abstractmethod + def on_remove_request(self, req_index) -> None: ... + + def prepare_extra_inputs(self, input_batch, req_states) -> dict: + return {} + + @abstractmethod + def postprocess(self, text_hidden, mm_outputs, input_batch, req_states): + return text_hidden, mm_outputs + + def dummy_run(self, state, num_tokens) -> None: + pass +``` + +未来 Talker 迁移将实现 `TalkerMTPPlugin`： +- `postprocess()` → 运行 `talker_mtp` 前向传播，将 `code_predictor_codes` 写入缓冲区 +- `prepare_extra_inputs()` → 注入预处理结果 +- `dummy_run()` → 捕获 talker CUDAGraph + +### 4.5 `OmniGPUModelRunner` — 薄继承层 + +对 v2 `GPUModelRunner` 的最小扩展。处理 `ModelState` 自身无法覆盖的 Omni 特定生命周期钩子。 + +```python +class OmniGPUModelRunner(GPUModelRunner): + """v2 GPUModelRunner 上的薄层，用于 Omni 生命周期钩子。""" + + def load_model(self, *args, **kwargs): + super().load_model(*args, **kwargs) + # 用 Omni 感知版本替换 model_state + self.model_state = init_omni_model_state( + self.vllm_config, self.model, self.encoder_cache, self.device + ) + + def finish_requests(self, scheduler_output): + # 在上游移除请求索引之前清理中间缓冲区 + for req_id in scheduler_output.finished_req_ids | scheduler_output.preempted_req_ids: + if req_id in self.req_states.req_id_to_index: + idx = self.req_states.req_id_to_index[req_id] + self.model_state.remove_request(idx) + super().finish_requests(scheduler_output) +``` + +**为何覆写 `finish_requests`？** v2 `ModelState` 接口没有 `remove_request()`。我们在 `OmniModelState` 中添加它并从 runner 调用。这是我们需要的唯一 runner 级生命周期钩子。如果上游后续将 `remove_request()` 加入 `ModelState`，我们可以移除此覆写。 + +### 4.6 `OmniARModelRunner` — Thinker 的 Runner + +扩展 `OmniGPUModelRunner`，添加 AR 特定行为：`OmniOutput` 后处理、每请求 `pooler_output` 构建和 KV 提取。 + +```python +class OmniARModelRunner(OmniGPUModelRunner): + """AR 阶段 runner。生成每请求的 hidden states + 多模态输出。""" + + def __init__(self, vllm_config, device): + super().__init__(vllm_config, device) + self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config( + self.vllm_config, self.model_config + ) + + @torch.inference_mode() + def execute_model(self, scheduler_output, intermediate_tensors=None, + dummy_run=False, skip_attn_for_dummy_run=False): + # 在状态清理之前处理已完成请求的 KV 传输 + self._handle_kv_transfer_pre(scheduler_output) + + # 委托给 v2 GPUModelRunner.execute_model（不修改） + result = super().execute_model( + scheduler_output, intermediate_tensors, + dummy_run=dummy_run, skip_attn_for_dummy_run=skip_attn_for_dummy_run, + ) + return result + + @torch.inference_mode() + def sample_tokens(self, grammar_output): + if self.execute_model_state is None: + return None + + input_batch = self.execute_model_state.input_batch + hidden_states = self.execute_model_state.hidden_states + kv_connector_output = self.execute_model_state.kv_connector_output + num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp + self.execute_model_state = None + + if not self.is_last_pp_rank: + sampled, num_sampled, num_rejected = pp_receive(...) + self.postprocess(input_batch, sampled, num_sampled, num_rejected) + return None + + # --- Omni 特有：后处理模型输出 --- + text_hidden, multimodal_outputs = self.model_state.postprocess_model_output( + hidden_states, input_batch, self.req_states + ) + + # --- 标准 v2 采样 --- + sampler_output, num_sampled, num_rejected = self.sample( + text_hidden, input_batch, grammar_output + ) + + # --- Omni 特有：构建 pooler_output --- + pooler_output = self._build_pooler_output( + text_hidden, multimodal_outputs, input_batch + ) + + # --- 后处理（Triton kernel 更新 req_states）--- + self.postprocess( + input_batch, sampler_output.sampled_token_ids, + num_sampled, num_rejected + ) + + output = OmniModelRunnerOutput( + req_ids=input_batch.req_ids, + req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)}, + sampled_token_ids=..., + pooler_output=pooler_output, + kv_connector_output=kv_connector_output, + kv_extracted_req_ids=self._kv_extracted_req_ids, + ) + + if self.use_async_scheduling: + return AsyncOutput(model_runner_output=output, ...) + return output + + def _build_pooler_output(self, text_hidden, mm_outputs, input_batch): + """按请求切分 hidden states + 多模态输出。""" + hidden_cpu = text_hidden.detach().cpu().contiguous() + mm_cpu = self._copy_mm_to_cpu(mm_outputs, hidden_cpu.shape[0]) + pooler = [] + for i, req_id in enumerate(input_batch.req_ids): + start = int(input_batch.query_start_loc_np[i]) + end = start + int(input_batch.num_scheduled_tokens[i]) + payload = {"hidden": hidden_cpu[start:end]} + payload.update(self._slice_mm_cpu(mm_cpu, start, end, i)) + pooler.append(payload) + return pooler + + def _handle_kv_transfer_pre(self, scheduler_output): + """在状态清理前为已完成请求提取 KV 缓存。""" + finished = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) + if finished and hasattr(self.model, "get_kv_transfer_metadata"): + for req_id, data in finished.items(): + meta = self.model.get_kv_transfer_metadata(req_id) + if meta: + existing = data.get("custom_metadata") or {} + existing.update(meta) + data["custom_metadata"] = existing + self._kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer( + finished_reqs=finished, + kv_caches=self.kv_caches, + block_size=self.cache_config.block_size, + cache_dtype=str(self.cache_config.cache_dtype), + ) +``` + +### 4.7 M-RoPE 集成 + +v2 的 `DefaultModelState` 已通过 `rope_state` 处理 M-RoPE： + +```python +# DefaultModelState.add_request() +if self.rope_state is not None: + self.rope_state.init_prefill_positions( + req_index, self.model, new_req_data.prefill_token_ids, + mm_features=new_req_data.mm_features, + ) + +# DefaultModelState.prepare_inputs() +positions = self.rope_state.get_positions(num_tokens_after_padding) +return {"positions": positions} +``` + +这会委托给模型的 `get_mrope_input_positions()` —— Qwen3-Omni Thinker 已经实现了该方法。**M-RoPE 无需额外工作。** v2 的 `rope_state` 处理： + +- 带多模态特征的 Prefill 位置初始化 +- Decode 位置更新（线性递增 + position delta） +- GPU 缓冲区的 Staged writes 管理 + +v1 的 `_init_mrope_positions` / `_calc_mrope_positions` / `_fixup_precomputed_mrope_decode_positions` 逻辑对于 Thinker 而言**完全被 `rope_state` 取代**。（GLM-Image 的预计算 decode 修正可能需要在未来阶段通过 `PositionEncodingStrategy` 覆写来实现。） + +### 4.8 `model_inputs` 注入流程 + +关键的 kwargs 注入通过 v2 现有的展开模式工作： + +```python +# v2 GPUModelRunner.execute_model() 第 991-998 行（不修改） +model_inputs = { + "input_ids": input_batch.input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": intermediate_tensors, + **self.model_state.prepare_inputs(input_batch, self.req_states), +} +``` + +`OmniModelState.prepare_inputs()` 返回： + +```python +{ + "positions": mrope_3d_positions, # 来自 rope_state + "model_intermediate_buffer": [dict, dict, ...], # 每请求 + "runtime_additional_information": [dict, ...], # 向后兼容别名 +} +``` + +这些键展开到 `model_inputs` 中，并传递给 `self.model(**model_inputs)`。Qwen3-Omni 统一的 `forward()` 接受 `runtime_additional_information` 作为显式参数，`**kwargs` 吸收其余部分： + +```python +# Qwen3OmniMoeForConditionalGeneration.forward() +def forward(self, input_ids, positions, ..., + runtime_additional_information=None, **kwargs): + if self.model_stage == "thinker": + # 不使用 runtime_additional_information 或 model_intermediate_buffer + text_hidden_states, captured_layer_dict = self.thinker( + input_ids=input_ids, positions=positions, + inputs_embeds=inputs_embeds, **capture_kwargs, **kwargs, + ) + return text_hidden_states, captured_layer_dict +``` + +**Thinker 静默吸收额外的 kwargs。** 无需修改模型代码。 + +### 4.9 `OmniOutput` 处理流水线 + +``` +model(**model_inputs) + │ + ▼ 返回 (text_hidden_states, captured_layer_dict) — 原始元组，不是 OmniOutput + │ + ▼ GPUModelRunner 将其作为 hidden_states 存入 ExecuteModelState + │ + ▼ OmniARModelRunner.sample_tokens() + │ + ▼ model_state.postprocess_model_output(hidden_states, ...) + │ ├── model.make_omni_output(raw_output, **buffer_kwargs) + │ │ └── 返回 OmniOutput(text_hidden_states=..., multimodal_outputs={ + │ │ captured layers, tts_bos_embed, tts_eos_embed, tts_pad_embed + │ │ }) + │ ├── 提取: text_hidden = omni_output.text_hidden_states + │ └── 提取: mm_outputs = omni_output.multimodal_outputs + │ + ▼ self.sample(text_hidden, input_batch, grammar_output) + │ └── compute_logits → sample → sampler_output + │ + ▼ _build_pooler_output(text_hidden, mm_outputs, input_batch) + │ └── 每请求: {hidden: slice, **mm_slice} + │ + ▼ OmniModelRunnerOutput(pooler_output=pooler, ...) +``` + +**重要边界情况：** v2 的 `execute_model` 期望最后一个 PP rank 的 `hidden_states` 为 `torch.Tensor`。但 Thinker 返回的是 `(Tensor, dict)` 元组。两种方案： + +1. **优选方案：** `OmniModelState.prepare_inputs()` 返回一个标志，使模型直接输出 `OmniOutput`（在 `forward` 内部调用 `make_omni_output`）。 +2. **兜底方案：** 在 `OmniGPUModelRunner` 中覆写 `execute_model` 拦截模型输出并包装。（约 10 行） + +我们采用方案 2 以确保安全，在 `self.model(**model_inputs)` 之后拦截： + +```python +# 在 OmniGPUModelRunner 中，仅覆写 execute_model 以包装输出 +def execute_model(self, scheduler_output, ...): + result = super().execute_model(scheduler_output, ...) + # 拦截：如果 execute_model_state 的 hidden_states 不是 Tensor， + # 说明模型返回了元组。进行包装和缓存。 + if self.execute_model_state is not None: + hs = self.execute_model_state.hidden_states + if not isinstance(hs, torch.Tensor) and hasattr(self.model, "make_omni_output"): + buffer_list = self.model_state.intermediate_buffer.gather( + self.execute_model_state.input_batch + ) + omni_output = self.model.make_omni_output( + hs, model_intermediate_buffer=buffer_list, + runtime_additional_information=buffer_list, + ) + # 用 OmniOutput 替换 hidden_states，供 sample_tokens 使用 + self.execute_model_state = self.execute_model_state._replace( + hidden_states=omni_output + ) + return result +``` + +### 4.10 功能开关共存 + +```python +# vllm_omni/worker/omni_ar_worker.py +import os + +class OmniARWorker: + def __init__(self, vllm_config, device): + use_v2 = os.environ.get("VLLM_OMNI_USE_V2_RUNNER", "0") == "1" + if use_v2: + from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner + self.model_runner = OmniARModelRunner(vllm_config, device) + else: + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + self.model_runner = GPUARModelRunner(vllm_config, device) +``` + +时间线： +1. **现在 → 4 月：** `VLLM_OMNI_USE_V2_RUNNER=1` 选择性启用 +2. **2026 年 5 月：** 翻转默认值为 `1`（v1 可通过 `=0` 使用） +3. **2026 年 6 月/7 月：** 移除 v1 runner + +--- + +## 5. 数据流：Thinker 完整路径 + +``` +SchedulerOutput + │ + ▼ +┌─ OmniARModelRunner.execute_model() ─────────────────────────────────┐ +│ 1. _handle_kv_transfer_pre(scheduler_output) │ +│ └── 为已完成请求提取 KV 缓存 │ +│ 2. super().execute_model() → v2 GPUModelRunner.execute_model() │ +│ ├── finish_requests() │ +│ │ └── model_state.remove_request() → 清理缓冲区 │ +│ ├── free_states() │ +│ ├── add_requests() │ +│ │ └── model_state.add_request() │ +│ │ ├── rope_state.init_prefill_positions() → M-RoPE 初始化 │ +│ │ └── intermediate_buffer.add_request() │ +│ ├── update_requests() → block_tables │ +│ ├── prepare_inputs() → InputBatch [Triton kernels] │ +│ ├── prepare_attn() → block_tables, slot_mappings │ +│ ├── model_state.prepare_attn() → attn_metadata │ +│ ├── model_state.get_mm_embeddings() → 视觉/音频编码 │ +│ ├── model_state.prepare_inputs() → { │ +│ │ positions (M-RoPE 3D), │ +│ │ model_intermediate_buffer, │ +│ │ runtime_additional_information │ +│ │ } │ +│ ├── model(**model_inputs) → (text_hidden, captured_layers) │ +│ └── ExecuteModelState 已存储 │ +│ 3. 后拦截: 包装输出 → 通过 make_omni_output 生成 OmniOutput │ +└──────────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─ OmniARModelRunner.sample_tokens() ─────────────────────────────────┐ +│ 1. model_state.postprocess_model_output() │ +│ ├── OmniOutput.text_hidden_states → text_hidden │ +│ └── OmniOutput.multimodal_outputs → {captured_layers, tts_*} │ +│ 2. self.sample(text_hidden, ...) → sampled_token_ids [Triton] │ +│ 3. _build_pooler_output() → 每请求 hidden + mm 切片 │ +│ 4. self.postprocess() → Triton kernel 更新 req_states │ +│ 5. → OmniModelRunnerOutput │ +│ {sampled_token_ids, pooler_output, kv_extracted_req_ids} │ +└──────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 6. 与 v1 的对比变化 + +| 方面 | v1（当前） | v2（本设计） | +|------|-----------|-------------| +| 基类 | `GPUModelRunner` v1（~6283 行） | `GPUModelRunner` v2（~1168 行） | +| `OmniGPUModelRunner` | ~1403 行，覆写 `_update_states`、`_preprocess`、`_model_forward`、`_dummy_run` | ~150 行，覆写 `load_model`、`finish_requests`、`execute_model`（10 行拦截） | +| `GPUARModelRunner` | ~677 行，覆写 `execute_model`、`sample_tokens` | `OmniARModelRunner` ~350 行，覆写 `execute_model`（前置钩子 + super）、`sample_tokens` | +| `model_intermediate_buffer` | `dict[str, dict]` 以 `req_id` 为键 | `list[dict]` 以 `req_index` 索引，通过 `OmniIntermediateBuffer` 管理 | +| M-RoPE | 自定义 `_init_mrope_positions`、`_calc_mrope_positions`、修正逻辑 | 完全委托给 v2 `rope_state` | +| `OmniOutput` 包装 | `_model_forward` 调用 `make_omni_output`，缓存于 `_omni_last_model_output` | `execute_model` 中后拦截，然后在 `sample_tokens` 中通过 `postprocess_model_output` 处理 | +| 输入准备 | Python 密集型 `_preprocess`（~180 行） | v2 Triton kernels + `ModelState.prepare_inputs()` | +| kwargs 注入 | `_model_forward` → `_build_model_kwargs_extra` | `prepare_inputs()` 返回字典，展开到 `model_inputs` | +| 请求清理 | `_update_states` 从字典中弹出 | `finish_requests` → `model_state.remove_request(idx)` | +| Omni runner 总行数 | ~2080（OmniGPU + AR） | ~500（OmniGPU + AR）— **精简 76%** | + +--- + +## 7. 健壮性：未来阶段兼容性 + +本设计已针对所有已知 Omni 模型阶段的需求进行验证： + +### 7.1 Talker（未来阶段） + +- **`has_preprocess = True`**：`OmniModelState.prepare_inputs()` 已包含 `model_intermediate_buffer`。Talker 的 `preprocess()` 在 `TalkerMTPPlugin.prepare_extra_inputs()` 钩子内按请求执行。 +- **`talker_mtp` 子模块**：`TalkerMTPPlugin.postprocess()` 在 `sample_tokens` 中运行 talker_mtp 前向传播。CUDAGraph 包装独立管理（与 v1 相同）。 +- **`prompt_embeds` 覆盖**：`OmniModelState` 可在 `prepare_inputs()` 中实现覆盖，或通过专用的 `PromptEmbedsPlugin` 实现。 +- **`gpu_resident_buffer_keys`**：`OmniIntermediateBuffer.update()` 遵循每模型的 GPU 常驻键设置。 + +### 7.2 TTS（未来阶段） + +- 与 Talker 相同，但 `TalkerMTPPlugin` 禁用 CUDAGraph（模型有内部 AR 循环）。 +- 通过 `TTSOmniModelState` 子类或插件上的配置标志实现。 + +### 7.3 Code2Wav / Generation 阶段（未来阶段） + +- 使用 `pool()` 路径（决策 #3）：`OmniGenerationModelRunner` 覆写 `pool()` 而非 `sample_tokens()`。 +- `runtime_additional_information` 用于获取 `left_context_size` — 已由 `prepare_inputs()` 注入。 +- `seq_token_counts` 通过插件 `prepare_extra_inputs()` 注入。 + +### 7.4 MiMoAudio、MammothModa2、GLM-Image（未来阶段） + +- **MiMoAudio**：`mm_features` 和 `req_id` 始终存在于缓冲区中（来自 RFC #1595 的泛化）。 +- **MammothModa2**：网格约束的 `generated_len` — 可在 `prepare_inputs()` 中从 `req_states` 计算。 +- **GLM-Image**：预计算的 M-RoPE decode 位置 — 需要 `PositionEncodingStrategy` 覆写（Thinker 不需要）。 + +--- + +## 8. 文件结构 + +``` +vllm_omni/worker_v2/ +├── __init__.py +├── omni_model_runner.py # OmniGPUModelRunner (~150 行) +├── omni_ar_model_runner.py # OmniARModelRunner (~350 行) +├── model_states/ +│ ├── __init__.py # init_omni_model_state() 工厂函数 +│ ├── omni_model_state.py # OmniModelState (通用基类, ~300 行) +│ ├── intermediate_buffer.py # OmniIntermediateBuffer (~120 行) +│ ├── plugin.py # OmniModelStatePlugin (ABC, ~30 行) +│ └── models/ # (当前为空, 未来: talker.py, tts.py) +│ └── __init__.py +└── outputs.py # 重导出 OmniModelRunnerOutput +``` + +过渡期间与 `vllm_omni/worker/` 共存。 + +--- + +## 9. 迁移检查清单 + +- [ ] **Phase 0：验证**（1 周） + - [ ] 确认 v2 `rope_state` 正确处理 Qwen3-Omni M-RoPE + - [ ] 确认 v2 `DefaultModelState.get_mm_embeddings()` 对 Qwen3 视觉/音频塔正常工作 + - [ ] 确认 `prepare_inputs()` 的 kwargs 无过滤地传递到 `model.forward()` + - [ ] 确认 `ExecuteModelState` 能容纳非 Tensor 的 `hidden_states`（元组） + +- [ ] **Phase 1：核心状态**（2 周） + - [ ] 实现 `OmniIntermediateBuffer` + - [ ] 实现 `OmniModelStatePlugin`（ABC） + - [ ] 实现 `OmniModelState`（继承 `DefaultModelState`） + - [ ] 实现 `init_omni_model_state()` 工厂函数 + - [ ] 编写缓冲区 add/remove/gather、state prepare_inputs 的单元测试 + +- [ ] **Phase 2：AR Runner**（2 周） + - [ ] 实现 `OmniGPUModelRunner`（薄继承） + - [ ] 实现 `OmniARModelRunner`（Thinker 路径） + - [ ] 在 `OmniARWorker` 中接入功能开关 + - [ ] 集成测试：纯文本生成 + - [ ] 集成测试：多模态输入（图像、音频、视频） + - [ ] 集成测试：M-RoPE 位置与 v1 一致 + - [ ] 集成测试：`pooler_output` 内容与 v1 一致 + - [ ] 集成测试：KV 传输 + +- [ ] **Phase 3：验收**（1 周） + - [ ] 端到端测试：Qwen3-Omni-30B-A3B-Instruct Thinker + - [ ] 性能基准：吞吐量、延迟、GPU 内存 vs. v1 + - [ ] 回归测试：所有现有 Thinker 阶段模型 + - [ ] 推测解码（EAGLE）兼容性检查 + +--- + +## 10. 单元测试规划 + +本节规划 `tests/worker_v2/` 目录下的单元测试，遵循现有 `tests/worker/` 的测试约定： +- 所有测试文件标记 `pytestmark = [pytest.mark.core_model, pytest.mark.cpu]` +- 通过 `object.__new__(RunnerClass)` 跳过真实构造函数，仅填充被测方法所需的最小属性 +- 使用 `monkeypatch` 替换模块级符号（如 `set_forward_context`） +- 使用简单的 Dummy/Stub 类替代真实的 `InputBatch`、`RequestState` 等 +- 不依赖 GPU、调度器或完整推理管线 + +### 10.1 文件结构 + +``` +tests/worker_v2/ +├── __init__.py +├── test_intermediate_buffer.py # OmniIntermediateBuffer 单元测试 +├── test_omni_model_state.py # OmniModelState 单元测试 +├── test_omni_model_state_plugin.py # OmniModelStatePlugin 接口测试 +├── test_omni_gpu_model_runner.py # OmniGPUModelRunner 单元测试 +├── test_omni_ar_model_runner.py # OmniARModelRunner 单元测试 +└── test_init_model_state.py # init_omni_model_state 工厂函数测试 +``` + +### 10.2 共用桩对象和工具函数 + +以下桩对象在多个测试文件中复用，可放置于 `tests/worker_v2/conftest.py` 或各文件内联： + +```python +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class DummyInputBatch: + """最小 InputBatch 桩，提供 idx_mapping_np、req_ids 等。""" + def __init__(self, req_ids, idx_mapping_np, num_tokens=4, + num_tokens_after_padding=4): + self.req_ids = list(req_ids) + self.idx_mapping_np = idx_mapping_np + self.num_reqs = len(req_ids) + self.num_tokens = num_tokens + self.num_tokens_after_padding = num_tokens_after_padding + self.query_start_loc_np = np.zeros(len(req_ids), dtype=np.int32) + self.num_scheduled_tokens = np.ones(len(req_ids), dtype=np.int32) + + +class DummyRequestState: + """最小 RequestState 桩。""" + def __init__(self): + self.req_id_to_index = {} + + +class DummyNewReqData: + """最小 NewRequestData 桩。""" + def __init__(self, req_id, prompt_embeds=None, + additional_information=None, mm_features=None, + prefill_token_ids=None): + self.req_id = req_id + self.prompt_embeds = prompt_embeds + self.additional_information = additional_information + self.mm_features = mm_features or [] + self.prefill_token_ids = prefill_token_ids or [1, 2, 3] + + +class DummySchedulerOutput: + """最小 SchedulerOutput 桩。""" + def __init__(self, finished_req_ids=None, preempted_req_ids=None): + self.finished_req_ids = set(finished_req_ids or []) + self.preempted_req_ids = set(preempted_req_ids or []) + + +class DummyModel(torch.nn.Module): + """可配置的模型桩。""" + def __init__(self, have_multimodal_outputs=True, + has_preprocess=False, has_postprocess=False): + super().__init__() + self.have_multimodal_outputs = have_multimodal_outputs + self.has_preprocess = has_preprocess + self.has_postprocess = has_postprocess + + def make_omni_output(self, raw, **kwargs): + from vllm_omni.model_executor.models.output_templates import OmniOutput + text_hidden, captured = raw + return OmniOutput( + text_hidden_states=text_hidden.reshape(-1, text_hidden.shape[-1]), + multimodal_outputs=captured, + ) +``` + +### 10.3 `test_intermediate_buffer.py` — OmniIntermediateBuffer + +测试中间缓冲区的核心生命周期：添加、移除、聚合、更新。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_add_request_stores_req_id` | 基本 `add_request` | `buffers[idx]["req_id"] == new_req_data.req_id` | +| `test_add_request_decodes_prompt_embeds` | 带 `prompt_embeds` 的请求 | `buffers[idx]["prompt_embeds_cpu"]` 是 CPU Tensor，形状正确 | +| `test_add_request_decodes_additional_information` | 带 `additional_information` 字典 | 字典键值正确写入 `buffers[idx]` | +| `test_add_request_stores_mm_features` | 带 `mm_features` 的请求 | `buffers[idx]["mm_features"]` 引用与输入一致 | +| `test_add_request_empty_payload` | `prompt_embeds=None`、`additional_information=None` | `buffers[idx]` 仅有 `req_id`，无多余键 | +| `test_remove_request_clears_slot` | `remove_request` 后检查 | `buffers[idx] == {}` | +| `test_remove_request_idempotent` | 连续两次 `remove_request` | 不抛异常，`buffers[idx] == {}` | +| `test_gather_returns_batch_order` | 3 个请求、乱序 `idx_mapping_np` | 返回列表顺序与 `idx_mapping_np` 一致 | +| `test_gather_empty_batch` | 空 `idx_mapping_np` | 返回空列表 | +| `test_gather_after_remove_returns_empty_dict` | 添加后移除再聚合 | 对应位置返回 `{}` | +| `test_update_merges_new_keys` | 调用 `update` 添加新键 | 新旧键共存 | +| `test_update_overwrites_existing_key` | 同键多次 `update` | 值为最后一次写入 | +| `test_update_detaches_tensors` | `update` 带 GPU-like Tensor | 存储的 Tensor 已 detach 且为 clone | +| `test_slot_reuse_after_remove` | 移除后在同一 index 添加新请求 | 新请求数据正确，无旧数据残留 | + +```python +# 示例测试 +def test_add_request_stores_req_id(): + buffer = OmniIntermediateBuffer(max_num_reqs=4) + req_data = DummyNewReqData(req_id="r1") + buffer.add_request(req_index=0, new_req_data=req_data) + assert buffer.buffers[0]["req_id"] == "r1" + +def test_gather_returns_batch_order(): + buffer = OmniIntermediateBuffer(max_num_reqs=4) + for i, rid in enumerate(["r0", "r1", "r2"]): + buffer.add_request(i, DummyNewReqData(req_id=rid)) + batch = DummyInputBatch( + req_ids=["r2", "r0"], + idx_mapping_np=np.array([2, 0], dtype=np.int32), + ) + result = buffer.gather(batch) + assert result[0]["req_id"] == "r2" + assert result[1]["req_id"] == "r0" +``` + +### 10.4 `test_omni_model_state.py` — OmniModelState + +测试 ModelState 的核心方法：请求生命周期、输入准备、输出后处理。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_init_detects_model_flags` | 构造时读取 `has_preprocess`、`have_multimodal_outputs` | 属性正确反映模型标志 | +| `test_add_request_delegates_to_super_and_buffer` | `add_request` 调用链 | `super().add_request` 和 `intermediate_buffer.add_request` 均被调用（monkeypatch 跟踪） | +| `test_remove_request_clears_buffer_and_plugins` | 注册一个 dummy plugin 后 `remove_request` | buffer 已清理，plugin 的 `on_remove_request` 被调用 | +| `test_prepare_inputs_includes_buffer_keys` | 无插件时 `prepare_inputs` | 返回字典包含 `model_intermediate_buffer` 和 `runtime_additional_information` | +| `test_prepare_inputs_merges_plugin_extras` | 注册返回 `{"extra_key": 42}` 的 plugin | 返回字典包含 `extra_key` | +| `test_prepare_inputs_buffer_alias_consistency` | 检查两个别名键 | `base["model_intermediate_buffer"] is base["runtime_additional_information"]` | +| `test_postprocess_wraps_tuple_via_make_omni_output` | 传入 `(Tensor, dict)` 元组 | 调用 `make_omni_output`，返回 `(text_hidden, mm_outputs)` | +| `test_postprocess_passes_through_plain_tensor` | 传入普通 Tensor | `text_hidden` 是原 Tensor，`mm_outputs == {}` | +| `test_postprocess_extracts_omni_output` | 传入 `OmniOutput` 对象 | 正确提取 `text_hidden_states` 和 `multimodal_outputs` | +| `test_postprocess_runs_plugin_chain` | 注册两个 plugin | 两个 plugin 的 `postprocess` 按顺序被调用 | +| `test_register_plugin_appends` | 多次 `register_plugin` | `self.plugins` 长度递增 | +| `test_prepare_dummy_inputs_includes_buffer` | `prepare_dummy_inputs` | 返回字典包含正确长度的 `model_intermediate_buffer` 列表 | + +```python +# 示例测试 +def test_prepare_inputs_includes_buffer_keys(): + state = _make_omni_model_state(model=DummyModel()) + input_batch = DummyInputBatch( + req_ids=["r1"], idx_mapping_np=np.array([0], dtype=np.int32) + ) + result = state.prepare_inputs(input_batch, DummyRequestState()) + assert "model_intermediate_buffer" in result + assert "runtime_additional_information" in result + assert isinstance(result["model_intermediate_buffer"], list) + +def test_postprocess_wraps_tuple_via_make_omni_output(): + model = DummyModel(have_multimodal_outputs=True) + state = _make_omni_model_state(model=model) + raw_output = (torch.randn(4, 8), {"layer_3": torch.randn(4, 8)}) + input_batch = DummyInputBatch( + req_ids=["r1"], idx_mapping_np=np.array([0], dtype=np.int32) + ) + text_hidden, mm = state.postprocess_model_output( + raw_output, input_batch, DummyRequestState() + ) + assert isinstance(text_hidden, torch.Tensor) + assert "layer_3" in mm +``` + +### 10.5 `test_omni_model_state_plugin.py` — 插件接口 + +验证插件协议的正确性，为 Talker 插件迁移提供信心。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_plugin_abc_cannot_instantiate` | 直接实例化 `OmniModelStatePlugin` | 抛出 `TypeError`（ABC） | +| `test_concrete_plugin_lifecycle` | 实现一个计数 plugin | `on_add_request` 和 `on_remove_request` 各被调用正确次数 | +| `test_plugin_prepare_extra_inputs_default_empty` | 默认实现 | 返回 `{}` | +| `test_plugin_postprocess_default_passthrough` | 默认实现 | 输入 `(text_hidden, mm)` 原样返回 | +| `test_multiple_plugins_compose` | 注册两个 plugin，一个修改 `mm_outputs` | 最终输出包含两个 plugin 的修改 | + +```python +# 示例测试 +class CountingPlugin(OmniModelStatePlugin): + def __init__(self): + self.add_count = 0 + self.remove_count = 0 + + def on_add_request(self, req_index, new_req_data): + self.add_count += 1 + + def on_remove_request(self, req_index): + self.remove_count += 1 + + def postprocess(self, text_hidden, mm_outputs, input_batch, req_states): + mm_outputs["plugin_marker"] = True + return text_hidden, mm_outputs + +def test_concrete_plugin_lifecycle(): + plugin = CountingPlugin() + state = _make_omni_model_state(model=DummyModel()) + state.register_plugin(plugin) + state.add_request(0, DummyNewReqData("r1")) + state.add_request(1, DummyNewReqData("r2")) + state.remove_request(0) + assert plugin.add_count == 2 + assert plugin.remove_count == 1 +``` + +### 10.6 `test_omni_gpu_model_runner.py` — OmniGPUModelRunner + +测试薄继承层的关键覆写。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_load_model_replaces_model_state` | `load_model` 后 `self.model_state` 类型 | `isinstance(runner.model_state, OmniModelState)` | +| `test_finish_requests_calls_remove_request` | 带 finished_req_ids 的 scheduler output | `model_state.remove_request` 按 req_index 被调用（monkeypatch 跟踪） | +| `test_finish_requests_handles_preempted` | preempted_req_ids 非空 | 预抢占的请求也触发 `remove_request` | +| `test_finish_requests_skips_unknown_req_ids` | finished_req_ids 中包含不存在的 req_id | 不抛异常，仅处理已知请求 | +| `test_execute_model_intercepts_tuple_output` | 模型返回 `(Tensor, dict)` 后检查 `execute_model_state` | `hidden_states` 被替换为 `OmniOutput` | +| `test_execute_model_passthrough_tensor_output` | 模型返回普通 Tensor | `hidden_states` 保持原样，不触发拦截 | + +```python +# 示例测试 +def test_finish_requests_calls_remove_request(monkeypatch): + runner = _make_v2_runner(req_ids=["r1", "r2"]) + remove_log = [] + monkeypatch.setattr( + runner.model_state, "remove_request", + lambda idx: remove_log.append(idx) + ) + scheduler_output = DummySchedulerOutput(finished_req_ids={"r1"}) + runner.finish_requests(scheduler_output) + assert len(remove_log) == 1 + assert remove_log[0] == runner.req_states.req_id_to_index["r1"] +``` + +### 10.7 `test_omni_ar_model_runner.py` — OmniARModelRunner + +测试 AR runner 的 `sample_tokens` 和 `pooler_output` 构建。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_sample_tokens_returns_none_when_no_state` | `execute_model_state = None` | 返回 `None` | +| `test_sample_tokens_calls_postprocess_model_output` | 正常流程 | `model_state.postprocess_model_output` 被调用（monkeypatch 跟踪） | +| `test_sample_tokens_builds_pooler_output` | 2 个请求 | `pooler_output` 长度为 2，每个包含 `hidden` 键 | +| `test_pooler_output_hidden_shape_matches_scheduled_tokens` | 请求各有不同 scheduled_tokens | 每个 `hidden` 切片长度与 `num_scheduled_tokens` 一致 | +| `test_pooler_output_includes_multimodal_slices` | `multimodal_outputs` 含 Tensor | `pooler_output` 中包含对应的多模态切片 | +| `test_pooler_output_includes_list_multimodal` | `multimodal_outputs` 含 list 类型 | 列表元素按 req_index 正确映射 | +| `test_pooler_output_empty_multimodal` | `multimodal_outputs = {}` | `pooler_output` 每项仅含 `hidden` | +| `test_pooler_output_is_none_for_text_engine` | `engine_output_type == "text"` | `pooler_output is None` | +| `test_handle_kv_transfer_pre_extracts_ids` | 带 `finished_requests_needing_kv_transfer` | `_kv_extracted_req_ids` 非空 | +| `test_handle_kv_transfer_pre_noop_when_empty` | 无需传输的请求 | `_kv_extracted_req_ids` 为空列表或 None | +| `test_sample_tokens_output_type` | 正常流程 | 返回类型为 `OmniModelRunnerOutput` | +| `test_sample_tokens_async_mode` | `use_async_scheduling = True` | 返回类型含 `AsyncOutput` 包装 | + +```python +# 示例测试 +def test_sample_tokens_builds_pooler_output(monkeypatch): + runner = _make_ar_runner(req_ids=["r1", "r2"], hidden_size=8) + text_hidden = torch.randn(4, 8) # 2 tokens per request + mm_outputs = {} + monkeypatch.setattr( + runner.model_state, "postprocess_model_output", + lambda *a, **kw: (text_hidden, mm_outputs) + ) + _setup_execute_model_state(runner, text_hidden) + output = OmniARModelRunner.sample_tokens(runner, grammar_output=None) + assert output is not None + assert len(output.pooler_output) == 2 + for item in output.pooler_output: + assert "hidden" in item + assert item["hidden"].shape[0] > 0 + +def test_pooler_output_hidden_shape_matches_scheduled_tokens(monkeypatch): + runner = _make_ar_runner(req_ids=["r1", "r2"], hidden_size=8) + # r1 调度 1 个 token，r2 调度 3 个 token + runner._input_batch_query_start_loc = np.array([0, 1], dtype=np.int32) + runner._input_batch_num_scheduled = np.array([1, 3], dtype=np.int32) + text_hidden = torch.randn(4, 8) + monkeypatch.setattr( + runner.model_state, "postprocess_model_output", + lambda *a, **kw: (text_hidden, {}) + ) + _setup_execute_model_state(runner, text_hidden) + output = OmniARModelRunner.sample_tokens(runner, grammar_output=None) + assert output.pooler_output[0]["hidden"].shape[0] == 1 + assert output.pooler_output[1]["hidden"].shape[0] == 3 +``` + +### 10.8 `test_init_model_state.py` — 工厂函数 + +测试架构分发逻辑。 + +| 测试函数 | 测试内容 | 验证要点 | +|----------|---------|---------| +| `test_omni_architecture_returns_omni_state` | 传入 Qwen3 架构名 | `isinstance(result, OmniModelState)` | +| `test_non_omni_architecture_falls_through` | 传入非 Omni 架构名 | `isinstance(result, DefaultModelState)` | +| `test_whisper_architecture_returns_whisper_state` | 传入 Whisper 架构名 | 不返回 `OmniModelState`（上游逻辑优先） | +| `test_multiple_architectures_with_omni` | `architectures` 列表含 Omni + 其他 | 命中 Omni 分支 | +| `test_empty_architectures_falls_through` | `architectures = []` | 返回 `DefaultModelState` | + +```python +# 示例测试 +def test_omni_architecture_returns_omni_state(): + vllm_config = _make_dummy_config( + architectures=["Qwen3OmniMoeForConditionalGeneration"] + ) + model = DummyModel() + result = init_omni_model_state(vllm_config, model, None, torch.device("cpu")) + assert isinstance(result, OmniModelState) + +def test_non_omni_architecture_falls_through(): + vllm_config = _make_dummy_config( + architectures=["LlamaForCausalLM"] + ) + model = DummyModel() + result = init_omni_model_state(vllm_config, model, None, torch.device("cpu")) + assert not isinstance(result, OmniModelState) +``` + +### 10.9 测试覆盖矩阵 + +| 被测组件 | 单元测试文件 | 覆盖的关键路径 | +|---------|------------|--------------| +| `OmniIntermediateBuffer` | `test_intermediate_buffer.py` | add / remove / gather / update / 槽位复用 / 空批次 | +| `OmniModelState` | `test_omni_model_state.py` | 请求生命周期 / prepare_inputs / postprocess / 插件分发 / dummy inputs | +| `OmniModelStatePlugin` | `test_omni_model_state_plugin.py` | ABC 约束 / 生命周期回调 / 默认实现 / 多插件组合 | +| `OmniGPUModelRunner` | `test_omni_gpu_model_runner.py` | load_model 替换 / finish_requests 清理 / execute_model 输出拦截 | +| `OmniARModelRunner` | `test_omni_ar_model_runner.py` | sample_tokens 全流程 / pooler_output 构建 / KV 传输前置 / 异步模式 | +| `init_omni_model_state` | `test_init_model_state.py` | 架构分发 / 向后兼容 / 边界情况 | + +### 10.10 不在单元测试范围内（需集成/E2E 测试） + +以下场景需要真实的模型权重、GPU 和/或调度器配合，不适合纯 CPU 单元测试： + +| 场景 | 测试层级 | 说明 | +|------|---------|------| +| M-RoPE 位置正确性 | 集成测试 | 需要真实的 `rope_state` + 多模态特征 | +| 多模态编码器执行 | 集成测试 | 需要 `EncoderRunner` + 实际视觉/音频塔 | +| CUDA Graph 捕获和回放 | GPU 集成测试 | 需要 GPU + `CUDAGraphManager` | +| 完整 `execute_model` → `sample_tokens` 链路 | 集成测试 | 需要真实的注意力后端、KV 缓存 | +| v1 vs v2 输出一致性 | 回归测试 | 需要同一模型在两个 runner 上运行对比 | +| 性能基准 | 基准测试 | 需要 GPU + 真实模型 + 多请求负载 | +| 端到端 Qwen3-Omni Thinker | E2E 测试 | 需要完整的模型权重和推理管线 | + +--- + +## 11. 风险与缓解措施 + +| 风险 | 影响 | 缓解措施 | +|------|------|---------| +| v2 `ExecuteModelState` 期望 `hidden_states: torch.Tensor`，但 Thinker 返回元组 | `sample_tokens` 运行时崩溃 | 在 `execute_model` 中后拦截，存储前包装为 `OmniOutput` | +| v2 异步 `AsyncOutput` 与 `pooler_output` CPU 拷贝竞争 | 数据损坏 | 在同一 `copy_stream` 中序列化 `pooler_output` 拷贝，或在 `postprocess()` 之前同步构建 | +| `prepare_inputs()` 返回字典的键被未来 v2 版本过滤 | kwargs 静默丢失 | 集成测试断言 kwargs 到达 `model.forward()`；兜底：覆写 `execute_model` | +| `rope_state` 不支持 Qwen3-Omni 的 `get_mrope_input_positions` 签名 | 位置错误 | Phase 0 验证；兜底：自定义 `RopeState` 子类 | +| v2 `model_runner.py` CUDA graph 路径绕过 `model_inputs` 字典 | graph 模式下缺少 kwargs | FULL graph 模式回放捕获的缓冲区；`prepare_inputs()` 在捕获期间被调用。Phase 0 中验证。 | + +--- + +## 12. 待定事项 + +1. **`prompt_logprobs` 支持：** v2 有 `PromptLogprobsWorker`。需验证其与 `OmniOutput` 的兼容性（text_hidden 提取必须在 logprobs 计算之前完成）。 +2. **推测解码：** v2 的 `speculator.propose()` 期望标准 `hidden_states`。需验证与后处理后的 `text_hidden` 的兼容性。 +3. **v2 中的 `has_preprocess` 路径：** Thinker 不需要，但设计应文档化 Talker 如何在没有 v1 `_preprocess` 循环的情况下处理每请求的 `preprocess()` 调用。候选方案：基于插件的 `prepare_extra_inputs()`，在插件内部进行每请求迭代。 +4. **`prompt_embeds` 的 `UvaBackedTensor`：** 需要性能分析以确定 UVA 访问模式是否优于显式异步拷贝。推迟到 Phase 1。 diff --git a/docs/user_guide/diffusion/multi_replica_serving.md b/docs/user_guide/diffusion/multi_replica_serving.md new file mode 100644 index 000000000..2da8c698f --- /dev/null +++ b/docs/user_guide/diffusion/multi_replica_serving.md @@ -0,0 +1,80 @@ +# Diffusion Multi-Replica Serving + +Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity. + +## When to Use + +Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster. + +## Configuration + +A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica: + +```yaml +stages: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3" + num_replicas: 4 + engine_args: + model: Qwen/Qwen-Image + parallel_config: + world_size: 1 +``` + +For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices: + +```yaml +stages: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3,4,5,6,7" + num_replicas: 4 + engine_args: + model: Wan-AI/Wan2.2-T2V-A14B-Diffusers + parallel_config: + world_size: 2 +``` + +The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`. + +## Serving + +Start the server with the stage config: + +```bash +vllm serve Qwen/Qwen-Image \ + --omni \ + --port 8099 \ + --stage-configs-path /path/to/diffusion_multi_replica.yaml +``` + +Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request. + +## Benchmarking + +Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running: + +```bash +python benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8099 \ + --model Qwen/Qwen-Image \ + --task t2i \ + --dataset random \ + --num-prompts 8 \ + --request-rate inf \ + --num-inference-steps 2 \ + --width 256 --height 256 \ + --output-file diffusion_replica_4.json +``` + +For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs. + +## Notes and Limitations + +- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas. +- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices. +- Request-level latency may stay similar or slightly increase because each request still runs on one replica. +- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory. diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index 7d051ac4a..439d71d09 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -257,8 +257,11 @@ async def run_single_request( new_chunks = audio_data[audio_list_consumed:] audio_list_consumed = len(audio_data) elif isinstance(audio_data, torch.Tensor): - new_chunks = [audio_data] - audio_last_tensor = audio_data + if audio_data is not audio_last_tensor: + new_chunks = [audio_data] + audio_last_tensor = audio_data + else: + new_chunks = [] else: new_chunks = [] diff --git a/examples/offline_inference/qwen3_omni/run_tests.sh b/examples/offline_inference/qwen3_omni/run_tests.sh new file mode 100755 index 000000000..5e216f381 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_tests.sh @@ -0,0 +1,140 @@ +#!/usr/bin/env bash +# ============================================================================= +# Qwen3-Omni end2end functional test script +# Uses default built-in assets, iterates over query-type x output-modalities combinations +# ============================================================================= + +set -euo pipefail + +export VLLM_OMNI_USE_V2_RUNNER=1 + +MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +END2END="$SCRIPT_DIR/end2end.py" +LOG_DIR="$SCRIPT_DIR/test_logs" +SUMMARY_FILE="$LOG_DIR/summary.log" + +rm -rf "$LOG_DIR" +mkdir -p "$LOG_DIR" +> "$SUMMARY_FILE" + +PASS=0 +FAIL=0 +TOTAL=0 + +run_test() { + local name="$1" + shift + local log_file="$LOG_DIR/${name}.log" + local output_dir="$LOG_DIR/${name}_output" + + TOTAL=$((TOTAL + 1)) + echo "========================================" + echo "[${TOTAL}] Running: ${name}" + echo " Command: python $END2END $*" + echo "========================================" + + local start_time + start_time=$(date +%s) + + if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => PASS (${elapsed}s)" + echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE" + PASS=$((PASS + 1)) + else + local exit_code=$? + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)" + echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE" + # Append last 30 lines of log to summary + echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE" + tail -n 30 "$log_file" >> "$SUMMARY_FILE" + echo " --- end ---" >> "$SUMMARY_FILE" + FAIL=$((FAIL + 1)) + fi + echo "" +} + +# ============================================================================= +# Test cases +# ============================================================================= + +# --- 1. Text-only input --- +run_test "text_to_text" \ + --model "$MODEL_PATH" --query-type text --modalities text + +run_test "text_to_text_audio" \ + --model "$MODEL_PATH" --query-type text --modalities text,audio + +run_test "text_to_audio" \ + --model "$MODEL_PATH" --query-type text --modalities audio + +# --- 2. Image input --- +run_test "image_to_text" \ + --model "$MODEL_PATH" --query-type use_image --modalities text + +run_test "image_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_image --modalities text,audio + +# --- 3. Audio input --- +run_test "audio_to_text" \ + --model "$MODEL_PATH" --query-type use_audio --modalities text + +run_test "audio_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_audio --modalities text,audio + +# --- 4. Video input --- +run_test "video_to_text" \ + --model "$MODEL_PATH" --query-type use_video --modalities text + +run_test "video_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_video --modalities text,audio + +# --- 5. Multi-audio input --- +run_test "multi_audios_to_text" \ + --model "$MODEL_PATH" --query-type use_multi_audios --modalities text + +run_test "multi_audios_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_multi_audios --modalities text,audio + +# --- 6. Mixed modality input (audio+image+video) --- +run_test "mixed_to_text" \ + --model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text + +run_test "mixed_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text,audio + +# --- 7. Audio from video --- +run_test "audio_in_video_to_text" \ + --model "$MODEL_PATH" --query-type use_audio_in_video --modalities text + +run_test "audio_in_video_to_text_audio" \ + --model "$MODEL_PATH" --query-type use_audio_in_video --modalities text,audio + +# --- 8. py-generator mode --- +run_test "text_to_text_generator" \ + --model "$MODEL_PATH" --query-type text --modalities text --py-generator + +run_test "audio_to_text_audio_generator" \ + --model "$MODEL_PATH" --query-type use_audio --modalities text,audio --py-generator + +# --- 9. Multi-prompt batch inference --- +run_test "text_to_text_batch3" \ + --model "$MODEL_PATH" --query-type text --modalities text --num-prompts 3 + +# ============================================================================= +# Summary +# ============================================================================= +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE" +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "" +echo "Detailed logs: $LOG_DIR" +echo "Summary report: $SUMMARY_FILE" + +exit $FAIL diff --git a/examples/offline_inference/qwen3_omni/run_tests_async_chunk.sh b/examples/offline_inference/qwen3_omni/run_tests_async_chunk.sh new file mode 100755 index 000000000..b33650edc --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_tests_async_chunk.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +# ============================================================================= +# Qwen3-Omni end2end_async_chunk functional test script +# Uses default built-in assets, iterates over AsyncOmni query-type x output-modalities combinations +# ============================================================================= + +set -euo pipefail + +export VLLM_OMNI_USE_V2_RUNNER=1 + +MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +END2END="$SCRIPT_DIR/end2end_async_chunk.py" +STAGE_CONFIGS="$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" +LOG_DIR="$SCRIPT_DIR/test_logs_async_chunk" +SUMMARY_FILE="$LOG_DIR/summary.log" + +rm -rf "$LOG_DIR" +mkdir -p "$LOG_DIR" +> "$SUMMARY_FILE" + +PASS=0 +FAIL=0 +TOTAL=0 + +run_test() { + local name="$1" + shift + local log_file="$LOG_DIR/${name}.log" + local output_dir="$LOG_DIR/${name}_output" + + TOTAL=$((TOTAL + 1)) + echo "========================================" + echo "[${TOTAL}] Running: ${name}" + echo " Command: python $END2END $*" + echo "========================================" + + local start_time + start_time=$(date +%s) + + if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => PASS (${elapsed}s)" + echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE" + PASS=$((PASS + 1)) + else + local exit_code=$? + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)" + echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE" + # Append last 30 lines of log to summary + echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE" + tail -n 30 "$log_file" >> "$SUMMARY_FILE" + echo " --- end ---" >> "$SUMMARY_FILE" + FAIL=$((FAIL + 1)) + fi + echo "" +} + +# ============================================================================= +# Test cases +# ============================================================================= + +# --- 1. Text-only input --- +run_test "text_to_text" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text + +run_test "text_to_text_audio" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text,audio + +run_test "text_to_audio" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities audio + +# --- 2. Image input --- +run_test "image_to_text" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_image --modalities text + +run_test "image_to_text_audio" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_image --modalities text,audio + +# --- 3. Audio input --- +run_test "audio_to_text" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_audio --modalities text + +run_test "audio_to_text_audio" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_audio --modalities text,audio + +# --- 4. Video input --- +run_test "video_to_text" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_video --modalities text + +run_test "video_to_text_audio" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_video --modalities text,audio + +# --- 5. Multi-prompt batch inference --- +run_test "text_to_text_batch3" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text --num-prompts 3 + +# --- 6. Concurrent requests (max-in-flight) --- +run_test "text_to_text_audio_concurrent2" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text,audio --num-prompts 2 --max-in-flight 2 + +# --- 7. stream-audio-to-disk mode --- +run_test "audio_to_text_audio_stream" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type use_audio --modalities text,audio --stream-audio-to-disk + +# --- 8. request-timeout test (long enough timeout, should PASS) --- +run_test "text_to_text_with_timeout" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text --request-timeout-s 600 + +# --- 9. batch-timeout test (long enough timeout, should PASS) --- +run_test "text_to_text_audio_batch_timeout" \ + --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \ + --query-type text --modalities text,audio --batch-timeout-s 600 + +# ============================================================================= +# Summary +# ============================================================================= +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE" +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "" +echo "Detailed logs: $LOG_DIR" +echo "Summary report: $SUMMARY_FILE" + +exit $FAIL diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 8150338d5..d33958d7f 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -135,17 +135,18 @@ def _estimate_prompt_len( return 2048 -def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: +def get_custom_voice_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult: """Build CustomVoice sample inputs. Args: use_batch_sample: When True, return a batch of prompts; otherwise a single prompt. + model_name: Override the default HF model path (e.g. with a local path). Returns: QueryResult with Omni inputs and the CustomVoice model path. """ task_type = "CustomVoice" - model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" if use_batch_sample: texts = [ "其实我真的有发现，我是一个特别善于观察别人情绪的人。", @@ -196,17 +197,18 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: ) -def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: +def get_voice_design_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult: """Build VoiceDesign sample inputs. Args: use_batch_sample: When True, return a batch of prompts; otherwise a single prompt. + model_name: Override the default HF model path (e.g. with a local path). Returns: QueryResult with Omni inputs and the VoiceDesign model path. """ task_type = "VoiceDesign" - model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" + model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" if use_batch_sample: texts = [ "哥哥，你回来啦，人家等了你好久好久了，要抱抱！", @@ -255,18 +257,19 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: ) -def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult: +def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl", model_name: str | None = None) -> QueryResult: """Build Base (voice clone) sample inputs. Args: use_batch_sample: When True, return a batch of prompts (Case 2). mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior. + model_name: Override the default HF model path (e.g. with a local path). Returns: QueryResult with Omni inputs and the Base model path. """ task_type = "Base" - model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" + model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-Base" ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav" ref_audio_single = ref_audio_path_1 ref_text_single = ( @@ -335,14 +338,17 @@ def _build_inputs(args) -> tuple[str, list]: ) query_func = query_map[args.query_type] + model_override = args.model if args.model else None if args.query_type in {"CustomVoice", "VoiceDesign"}: - query_result = query_func(use_batch_sample=args.use_batch_sample) + query_result = query_func(use_batch_sample=args.use_batch_sample, model_name=model_override) elif args.query_type == "Base": - query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag) + query_result = query_func( + use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag, model_name=model_override + ) else: query_result = query_func() - model_name = query_result.model_name + model_name = args.model if args.model else query_result.model_name if args.txt_prompts: with open(args.txt_prompts) as f: @@ -359,7 +365,10 @@ def _build_inputs(args) -> tuple[str, list]: for t in lines ] else: - inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs] + if isinstance(query_result.inputs, list): + inputs = query_result.inputs * args.num_prompts + else: + inputs = [query_result.inputs] * args.num_prompts return model_name, inputs @@ -434,6 +443,12 @@ async def main_streaming(args): def parse_args(): parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser.add_argument( + "--model", + type=str, + default=None, + help="Override the default model path for the selected query type.", + ) parser.add_argument( "--query-type", "-q", diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end_benchmark_script.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end_benchmark_script.py new file mode 100644 index 000000000..066097e55 --- /dev/null +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end_benchmark_script.py @@ -0,0 +1,640 @@ +"""Offline benchmark script for Qwen3 TTS via vLLM Omni. + +Supports warmup + test rounds with a single Omni instance to avoid +repeated model loading. Prints per-round wall time summary at the end. +""" + +import asyncio +import logging +import os +import time +from typing import Any, NamedTuple + +import soundfile as sf +import torch + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +logger = logging.getLogger(__name__) + + +class QueryResult(NamedTuple): + """Container for a prepared Omni request.""" + + inputs: dict + model_name: str + + +_model_cache: dict[str, Any] = {} + + +def _estimate_prompt_len( + additional_information: dict[str, Any], + model_name: str, +) -> int: + """Estimate prompt_token_ids placeholder length for the Talker stage.""" + try: + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, + ) + + if model_name not in _model_cache: + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained( + model_name, trust_remote_code=True, padding_side="left", fix_mistral_regex=True + ) + cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True) + + speech_tok = None + try: + from transformers.utils import cached_file + + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer + + st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json") + if st_cfg_path: + speech_tok = Qwen3TTSTokenizer.from_pretrained( + # os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16 + os.path.dirname(st_cfg_path), + torch_dtype=torch.bfloat16, + fix_mistral_regex=True, + ) + logger.info("Loaded speech tokenizer for exact ref_code_len estimation") + except Exception as e: + logger.debug("Could not load speech tokenizer: %s", e) + + _model_cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok) + + tok, tcfg, speech_tok = _model_cache[model_name] + task_type = (additional_information.get("task_type") or ["CustomVoice"])[0] + + def _estimate_ref_code_len(ref_audio: object) -> int | None: + if not isinstance(ref_audio, (str, list)): + return None + audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio + if not isinstance(audio_path, str) or not audio_path.strip(): + return None + try: + from vllm.multimodal.media import MediaConnector + + connector = MediaConnector(allowed_local_media_path="/") + audio, sr = connector.fetch_audio(audio_path) + import numpy as np + + wav_np = np.asarray(audio, dtype=np.float32) + + if speech_tok is not None: + enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True) + ref_code = getattr(enc, "audio_codes", None) + if isinstance(ref_code, list): + ref_code = ref_code[0] if ref_code else None + if ref_code is not None and hasattr(ref_code, "shape"): + shape = ref_code.shape + return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None + + codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12 + return int(len(audio) / sr * codec_hz) + except Exception: + return None + + return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=additional_information, + task_type=task_type, + tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], + codec_language_id=getattr(tcfg, "codec_language_id", None), + spk_is_dialect=getattr(tcfg, "spk_is_dialect", None), + estimate_ref_code_len=_estimate_ref_code_len, + ) + except Exception as exc: + logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc) + return 2048 + + +def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: + task_type = "CustomVoice" + model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + if use_batch_sample: + texts = [ + "其实我真的有发现，我是一个特别善于观察别人情绪的人。", + "She said she would be here by noon.", + "I like you very much.", + "Really, you do?", + "Yes, absolutely.", + ] + instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."] + languages = ["Chinese", "English", "English", "English", "English"] + speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"] + inputs = [] + for text, instruct, language, speaker in zip(texts, instructs, languages, speakers): + additional_information = { + "task_type": [task_type], + "text": [text], + "instruct": [instruct], + "language": [language], + "speaker": [speaker], + "max_new_tokens": [2048], + } + inputs.append( + { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + ) + else: + text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。" + language = "Chinese" + speaker = "Vivian" + instruct = "用特别愤怒的语气说" + additional_information = { + "task_type": [task_type], + "text": [text], + "language": [language], + "speaker": [speaker], + "instruct": [instruct], + "max_new_tokens": [2048], + } + inputs = { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + return QueryResult(inputs=inputs, model_name=model_name) + + +def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: + task_type = "VoiceDesign" + model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" + if use_batch_sample: + texts = [ + "哥哥，你回来啦，人家等了你好久好久了，要抱抱！", + "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!", + ] + instructs = [ + "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。", + "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.", + ] + languages = ["Chinese", "English"] + inputs = [] + for text, instruct, language in zip(texts, instructs, languages): + additional_information = { + "task_type": [task_type], + "text": [text], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [2048], + "non_streaming_mode": [True], + } + inputs.append( + { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + ) + else: + text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！" + instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。" + language = "Chinese" + additional_information = { + "task_type": [task_type], + "text": [text], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [2048], + "non_streaming_mode": [True], + } + inputs = { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + return QueryResult(inputs=inputs, model_name=model_name) + + +def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult: + task_type = "Base" + model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" + ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav" + ref_audio_single = ref_audio_path_1 + ref_text_single = ( + "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." + ) + syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye." + syn_lang_single = "Auto" + x_vector_only_mode = mode_tag == "xvec_only" + if use_batch_sample: + syn_text_batch = [ + "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.", + "其实我真的有发现，我是一个特别善于观察别人情绪的人。", + ] + syn_lang_batch = ["Chinese", "English"] + inputs = [] + for text, language in zip(syn_text_batch, syn_lang_batch): + additional_information = { + "task_type": [task_type], + "ref_audio": [ref_audio_single], + "ref_text": [ref_text_single], + "text": [text], + "language": [language], + "x_vector_only_mode": [x_vector_only_mode], + "max_new_tokens": [2048], + } + inputs.append( + { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + ) + else: + additional_information = { + "task_type": [task_type], + "ref_audio": [ref_audio_single], + "ref_text": [ref_text_single], + "text": [syn_text_single], + "language": [syn_lang_single], + "x_vector_only_mode": [x_vector_only_mode], + "max_new_tokens": [2048], + } + inputs = { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name), + "additional_information": additional_information, + } + return QueryResult(inputs=inputs, model_name=model_name) + + +query_map = { + "CustomVoice": get_custom_voice_query, + "VoiceDesign": get_voice_design_query, + "Base": get_base_query, +} + + +def _build_inputs(args) -> tuple[str, list]: + """Resolve model name and inputs list from CLI args.""" + if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0: + raise ValueError( + f"--batch-size must be a power of two (got {args.batch_size}); " + "non-power-of-two values do not align with CUDA graph capture sizes " + "of Code2Wav." + ) + + query_func = query_map[args.query_type] + if args.query_type in {"CustomVoice", "VoiceDesign"}: + query_result = query_func(use_batch_sample=args.use_batch_sample) + elif args.query_type == "Base": + query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag) + else: + query_result = query_func() + + model_name = query_result.model_name + + if args.txt_prompts: + with open(args.txt_prompts) as f: + lines = [line.strip() for line in f if line.strip()] + if not lines: + raise ValueError(f"No valid prompts found in {args.txt_prompts}") + template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0] + template_info = template["additional_information"] + inputs = [ + { + "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name), + "additional_information": {**template_info, "text": [t]}, + } + for t in lines + ] + else: + inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs] + + return model_name, inputs + + +def _save_wav(output_dir: str, request_id: str, mm: dict) -> None: + """Concatenate audio chunks and write to a wav file.""" + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV") + logger.info("Saved audio to %s", out_wav) + + +class _SummaryCollector: + """Monkey-patches omni._log_summary_and_cleanup to capture summary dicts. + + NOTE: This relies on a private method ``_log_summary_and_cleanup`` which may + break if the upstream Omni implementation is refactored. If Omni exposes a + public hook for summary collection in the future, prefer that instead. + """ + + def __init__(self, omni): + self.omni = omni + self.summaries: list[dict[str, Any]] = [] + self.collecting = False + self._orig = omni._log_summary_and_cleanup + + def install(self): + from pprint import pformat as _pformat + + collector = self + + def _patched(request_id: str) -> None: + req_state = collector.omni.request_states.get(request_id) + try: + if req_state is None or req_state.metrics is None: + return + summary = req_state.metrics.build_and_log_summary() + logger.info("[Summary] %s", _pformat(summary, sort_dicts=False)) + if collector.collecting and summary: + collector.summaries.append(summary) + except Exception: + logger.exception("Failed to build/log summary for req=%s", request_id) + finally: + collector.omni.request_states.pop(request_id, None) + + collector.omni._log_summary_and_cleanup = _patched + + +def _print_summary(round_times: list[float], summaries: list[dict[str, Any]] | None = None) -> None: + """Print per-round wall time and optional detailed pipeline summary.""" + print("\n" + "=" * 60) + print("BENCHMARK SUMMARY") + print("=" * 60) + for i, t in enumerate(round_times): + print(f" Test round {i}: {t:.3f}s") + print("-" * 60) + avg = sum(round_times) / len(round_times) + print(f" Rounds: {len(round_times)}") + print(f" Total: {sum(round_times):.3f}s") + print(f" Average: {avg:.3f}s") + print(f" Min: {min(round_times):.3f}s") + print(f" Max: {max(round_times):.3f}s") + print("=" * 60) + + if not summaries: + print() + return + + # --- Aggregate from captured summaries --- + overall_list = [s["overall_summary"] for s in summaries if "overall_summary" in s] + if not overall_list: + print() + return + + n = len(overall_list) + + # Average Wall time + avg_wall = sum(d["e2e_wall_time_ms"] for d in overall_list) / n + print(f"\n{'=' * 60}") + print("DETAILED PIPELINE SUMMARY (test rounds average)") + print("=" * 60) + print(f" Avg Wall Time: {avg_wall:.3f} ms") + + # Per-stage timing + stage_keys = sorted(k for k in overall_list[0] if k.startswith("e2e_stage_") and k.endswith("_wall_time_ms")) + if stage_keys: + print("-" * 60) + print(" Per-stage timing (avg):") + for sk in stage_keys: + stage_avg = sum(d.get(sk, 0.0) for d in overall_list) / n + # extract stage index from key like e2e_stage_0_wall_time_ms + stage_idx = sk.replace("e2e_stage_", "").replace("_wall_time_ms", "") + print(f" Stage {stage_idx} wall time: {stage_avg:.3f} ms") + + # Per-stage gen time from stage_table + all_stage_gen: dict[int, list[float]] = {} + for s in summaries: + for req_entry in s.get("stage_table", []): + for stage_info in req_entry.get("stages", []): + sid = stage_info.get("stage_id", -1) + gen_ms = stage_info.get("stage_gen_time_ms", 0.0) + all_stage_gen.setdefault(sid, []).append(gen_ms) + if all_stage_gen: + print(" Per-stage gen time (avg):") + for sid in sorted(all_stage_gen): + vals = all_stage_gen[sid] + print(f" Stage {sid} gen time: {sum(vals) / len(vals):.3f} ms") + + # Per-prompt throughput + avg_tokens_per_s = sum(d.get("e2e_avg_tokens_per_s", 0.0) for d in overall_list) / n + avg_time_per_req = sum(d.get("e2e_avg_time_per_request_ms", 0.0) for d in overall_list) / n + total_tokens_list = [d.get("e2e_total_tokens", 0) for d in overall_list] + avg_total_tokens = sum(total_tokens_list) / n + print("-" * 60) + print(" Per-prompt throughput (avg):") + print(f" Tokens per request: {avg_total_tokens:.1f}") + print(f" Time per request: {avg_time_per_req:.3f} ms") + print(f" Tokens/s: {avg_tokens_per_s:.3f}") + print("=" * 60 + "\n") + + +def main(args): + """Run sync benchmark: warmup rounds then test rounds.""" + model_name, inputs = _build_inputs(args) + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + omni = Omni( + model=model_name, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + ) + + collector = _SummaryCollector(omni) + collector.install() + + batch_size = args.batch_size + + # --- Warmup --- + for r in range(args.warmup_rounds): + print(f"[Warmup {r}/{args.warmup_rounds}] running...") + for batch_start in range(0, len(inputs), batch_size): + batch = inputs[batch_start : batch_start + batch_size] + for _ in omni.generate(batch): + pass + print(f"[Warmup {r}/{args.warmup_rounds}] done") + + # --- Test --- + collector.collecting = True + round_times = [] + for r in range(args.test_rounds): + print(f"[Test {r}/{args.test_rounds}] running...") + t_start = time.perf_counter() + for batch_start in range(0, len(inputs), batch_size): + batch = inputs[batch_start : batch_start + batch_size] + for stage_outputs in omni.generate(batch): + output = stage_outputs.request_output + _save_wav( + output_dir, + f"round{r}_{output.request_id}", + output.outputs[0].multimodal_output, + ) + elapsed = time.perf_counter() - t_start + round_times.append(elapsed) + print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s") + collector.collecting = False + + _print_summary(round_times, collector.summaries if args.log_stats else None) + + +async def main_streaming(args): + """Run async streaming benchmark: warmup rounds then test rounds.""" + model_name, inputs = _build_inputs(args) + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + omni = AsyncOmni( + model=model_name, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + ) + + collector = _SummaryCollector(omni) + collector.install() + + # --- Warmup --- + for r in range(args.warmup_rounds): + print(f"[Warmup {r}/{args.warmup_rounds}] running...") + for i, prompt in enumerate(inputs): + async for _ in omni.generate(prompt, request_id=f"warmup_{r}_{i}"): + pass + print(f"[Warmup {r}/{args.warmup_rounds}] done") + + # --- Test --- + collector.collecting = True + round_times = [] + for r in range(args.test_rounds): + print(f"[Test {r}/{args.test_rounds}] running...") + t_start = time.perf_counter() + for i, prompt in enumerate(inputs): + request_id = f"round{r}_{i}" + t_req_start = time.perf_counter() + t_prev = t_req_start + chunk_idx = 0 + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = stage_output.request_output.outputs[0].multimodal_output + if not stage_output.finished: + t_now = time.perf_counter() + audio = mm.get("audio") + n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + if chunk_idx == 0: + ttfa_ms = (t_now - t_req_start) * 1000 + logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms) + else: + dt_ms = (t_now - t_prev) * 1000 + logger.info( + "Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms + ) + t_prev = t_now + chunk_idx += 1 + else: + t_end = time.perf_counter() + total_ms = (t_end - t_req_start) * 1000 + logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx) + _save_wav(output_dir, request_id, mm) + elapsed = time.perf_counter() - t_start + round_times.append(elapsed) + print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s") + collector.collecting = False + + _print_summary(round_times, collector.summaries if args.log_stats else None) + + +def parse_args(): + parser = FlexibleArgumentParser(description="Benchmark script for Qwen3 TTS via vLLM Omni") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="CustomVoice", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--warmup-rounds", + type=int, + default=3, + help="Number of warmup rounds (results discarded, default: 3).", + ) + parser.add_argument( + "--test-rounds", + type=int, + default=7, + help="Number of test rounds (results saved and timed, default: 7).", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics (default: disabled).", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=300, + help="Timeout for initializing a single stage in seconds (default: 300).", + ) + parser.add_argument( + "--output-dir", + default="output_audio", + help="Output directory for generated wav files (default: output_audio).", + ) + parser.add_argument( + "--txt-prompts", + type=str, + default=None, + help="Path to a .txt file with one prompt per line.", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to a stage configs file.", + ) + parser.add_argument( + "--use-batch-sample", + action="store_true", + default=False, + help="Use batch input sample for CustomVoice/VoiceDesign/Base query.", + ) + parser.add_argument( + "--mode-tag", + type=str, + default="icl", + choices=["icl", "xvec_only"], + help="Mode tag for Base query x_vector_only_mode (default: icl).", + ) + parser.add_argument( + "--streaming", + action="store_true", + default=False, + help="Stream audio chunks via AsyncOmni.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Number of prompts per batch (default: 1).", + ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + if args.streaming: + asyncio.run(main_streaming(args)) + else: + main(args) diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/run_tests.sh b/examples/offline_inference/text_to_speech/qwen3_tts/run_tests.sh new file mode 100755 index 000000000..68389d0f3 --- /dev/null +++ b/examples/offline_inference/text_to_speech/qwen3_tts/run_tests.sh @@ -0,0 +1,173 @@ +#!/usr/bin/env bash +# ============================================================================= +# Qwen3-TTS end2end functional test script +# Iterates over query-type and mode combinations for qwen3_tts/end2end.py +# +# Usage: +# ./run_tests.sh --custom-voice-model /path/a --voice-design-model /path/b --base-model /path/c +# ./run_tests.sh --custom-voice-model /path/a # only run CustomVoice tests +# ./run_tests.sh --base-model /path/c # only run Base tests +# ============================================================================= + +set -euo pipefail + +export VLLM_OMNI_USE_V2_RUNNER=1 +# export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + +# ---- Parse named arguments ---- +CUSTOM_VOICE_MODEL="" +VOICE_DESIGN_MODEL="" +BASE_MODEL="" + +usage() { + echo "Usage: $0 [--custom-voice-model PATH] [--voice-design-model PATH] [--base-model PATH]" + echo "" + echo "Pass one or more model paths. Only the query types with a model provided will be tested." + exit 1 +} + +if [[ $# -eq 0 ]]; then + usage +fi + +while [[ $# -gt 0 ]]; do + case "$1" in + --custom-voice-model) CUSTOM_VOICE_MODEL="$2"; shift 2 ;; + --voice-design-model) VOICE_DESIGN_MODEL="$2"; shift 2 ;; + --base-model) BASE_MODEL="$2"; shift 2 ;; + -h|--help) usage ;; + *) echo "Unknown option: $1"; usage ;; + esac +done + +if [[ -z "$CUSTOM_VOICE_MODEL" && -z "$VOICE_DESIGN_MODEL" && -z "$BASE_MODEL" ]]; then + echo "Error: at least one model path must be provided." + usage +fi + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +END2END="$SCRIPT_DIR/end2end.py" +LOG_DIR="$SCRIPT_DIR/test_logs" +SUMMARY_FILE="$LOG_DIR/summary.log" + +rm -rf "$LOG_DIR" +mkdir -p "$LOG_DIR" +> "$SUMMARY_FILE" + +PASS=0 +FAIL=0 +TOTAL=0 + +run_test() { + local name="$1" + shift + local log_file="$LOG_DIR/${name}.log" + local output_dir="$LOG_DIR/${name}_output" + + TOTAL=$((TOTAL + 1)) + echo "========================================" + echo "[${TOTAL}] Running: ${name}" + echo " Command: python $END2END $*" + echo "========================================" + + local start_time + start_time=$(date +%s) + + if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => PASS (${elapsed}s)" + echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE" + PASS=$((PASS + 1)) + else + local exit_code=$? + local end_time + end_time=$(date +%s) + local elapsed=$((end_time - start_time)) + echo " => FAIL (exit code: ${exit_code}, ${elapsed}s)" + echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE" + echo " --- last 30 lines of log ---" >> "$SUMMARY_FILE" + tail -n 30 "$log_file" >> "$SUMMARY_FILE" + echo " --- end ---" >> "$SUMMARY_FILE" + FAIL=$((FAIL + 1)) + fi + echo "" +} + +# ============================================================================= +# Test cases — only run sections whose model path was provided +# ============================================================================= + +# --- 1. CustomVoice --- +if [[ -n "$CUSTOM_VOICE_MODEL" ]]; then + run_test "custom_voice" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice + + run_test "custom_voice_batch_sample" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --use-batch-sample + + run_test "custom_voice_py_generator" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --py-generator + + run_test "custom_voice_streaming" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --streaming + + run_test "custom_voice_batch4" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4 + + run_test "custom_voice_batch4_batch_sample" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4 --use-batch-sample --stage-configs-path "$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml" + + run_test "custom_voice_num_prompts3" \ + --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --num-prompts 3 +fi + +# --- 2. VoiceDesign --- +if [[ -n "$VOICE_DESIGN_MODEL" ]]; then + run_test "voice_design" \ + --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign + + run_test "voice_design_batch_sample" \ + --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --use-batch-sample + + run_test "voice_design_py_generator" \ + --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --py-generator + + run_test "voice_design_streaming" \ + --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --streaming +fi + +# --- 3. Base (icl mode) --- +if [[ -n "$BASE_MODEL" ]]; then + run_test "base_icl" \ + --model "$BASE_MODEL" --query-type Base --mode-tag icl + + run_test "base_icl_batch_sample" \ + --model "$BASE_MODEL" --query-type Base --mode-tag icl --use-batch-sample + + run_test "base_icl_py_generator" \ + --model "$BASE_MODEL" --query-type Base --mode-tag icl --py-generator + + run_test "base_icl_streaming" \ + --model "$BASE_MODEL" --query-type Base --mode-tag icl --streaming + + # --- 4. Base (xvec_only mode) --- + run_test "base_xvec_only" \ + --model "$BASE_MODEL" --query-type Base --mode-tag xvec_only + + run_test "base_xvec_only_batch_sample" \ + --model "$BASE_MODEL" --query-type Base --mode-tag xvec_only --use-batch-sample +fi + +# ============================================================================= +# Summary +# ============================================================================= +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE" +echo "========================================" | tee -a "$SUMMARY_FILE" +echo "" +echo "Detailed logs: $LOG_DIR" +echo "Summary report: $SUMMARY_FILE" + +exit $FAIL diff --git a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh index ed09bac14..c6293e4e2 100644 --- a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh +++ b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh @@ -1,6 +1,9 @@ #!/usr/bin/env bash set -euo pipefail +# Default model name +MODEL="${MODEL:-Qwen/Qwen2.5-Omni-7B}" + # Default query type QUERY_TYPE="${1:-mixed_modalities}" @@ -162,7 +165,7 @@ echo "" request_body=$(cat <= len(request.prompt_token_ids) + _adapter_done = ( + chunk_transfer_adapter is not None and request.request_id in chunk_transfer_adapter.finished_requests + ) + return ( + request.status == RequestStatus.FINISHED_STOPPED + or (_all_computed and chunk_transfer_adapter is None) + or (_all_computed and _adapter_done) + ) + + +class TestFinishConditionSyncMode(unittest.TestCase): + """Case 1: No adapter (sync mode) — finish when all tokens computed.""" + + def test_finishes_when_all_computed(self): + req = FakeRequest("r1", prompt_len=100, num_computed=100) + assert _evaluate_finish_condition(req, chunk_transfer_adapter=None) + + def test_does_not_finish_when_partially_computed(self): + req = FakeRequest("r1", prompt_len=100, num_computed=50) + assert not _evaluate_finish_condition(req, chunk_transfer_adapter=None) + + +class TestFinishConditionAsyncAdapterDone(unittest.TestCase): + """Case 2: Adapter present AND adapter signals done — finish.""" + + def test_finishes_when_computed_and_adapter_done(self): + req = FakeRequest("r1", prompt_len=100, num_computed=100) + adapter = FakeChunkAdapter(finished_request_ids={"r1"}) + assert _evaluate_finish_condition(req, adapter) + + def test_does_not_finish_when_computed_but_adapter_not_done(self): + """Case 3: Adapter present but NOT done — must NOT finish.""" + req = FakeRequest("r1", prompt_len=100, num_computed=100) + adapter = FakeChunkAdapter(finished_request_ids=set()) + assert not _evaluate_finish_condition(req, adapter) + + +class TestFinishConditionAsyncNotComputed(unittest.TestCase): + """Adapter present, tokens not all computed — never finishes.""" + + def test_does_not_finish_partial_compute_adapter_done(self): + req = FakeRequest("r1", prompt_len=100, num_computed=50) + adapter = FakeChunkAdapter(finished_request_ids={"r1"}) + assert not _evaluate_finish_condition(req, adapter) + + def test_does_not_finish_partial_compute_adapter_not_done(self): + req = FakeRequest("r1", prompt_len=100, num_computed=50) + adapter = FakeChunkAdapter(finished_request_ids=set()) + assert not _evaluate_finish_condition(req, adapter) + + +class TestFinishConditionAlreadyStopped(unittest.TestCase): + """Request already FINISHED_STOPPED — always returns True.""" + + def test_already_stopped_no_adapter(self): + req = FakeRequest("r1", prompt_len=100, num_computed=0) + req.status = RequestStatus.FINISHED_STOPPED + assert _evaluate_finish_condition(req, chunk_transfer_adapter=None) + + def test_already_stopped_with_adapter(self): + req = FakeRequest("r1", prompt_len=100, num_computed=0) + req.status = RequestStatus.FINISHED_STOPPED + adapter = FakeChunkAdapter(finished_request_ids=set()) + assert _evaluate_finish_condition(req, adapter) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py index 121551664..c98aeb5d9 100644 --- a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py +++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py @@ -61,6 +61,7 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] pytestmark = [pytest.mark.full_model, pytest.mark.omni] _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") +_CI_MIN_DAILY_OMNI_ACCURACY = "0.685" def get_chunk_config(config_path: str | None = None): @@ -119,9 +120,11 @@ def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None: _require_vllm_cli() pytest.importorskip("datasets") pytest.importorskip("huggingface_hub") - ns = _acc_bench.parse_acc_benchmark_args( - build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False) - ) + argv = build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False) + # Daily-Omni's full-set score can move by a few MCQ items across GPU and + # batching schedules; keep CI below the nominal 0.69 CLI baseline. + argv.extend(["--min-daily-omni-accuracy", _CI_MIN_DAILY_OMNI_ACCURACY]) + ns = _acc_bench.parse_acc_benchmark_args(argv) assert _acc_bench.run_acc_benchmark(ns) == 0 diff --git a/tests/e2e/offline_inference/test_audiox_model.py b/tests/e2e/offline_inference/test_audiox_model.py new file mode 100644 index 000000000..f353b2451 --- /dev/null +++ b/tests/e2e/offline_inference/test_audiox_model.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +# Prefer a tiny/random checkpoint for CI. +# Override in CI if needed: AUDIOX_TEST_MODEL= +models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4", "xpu": "B60"}) +@pytest.mark.parametrize("model_name", models) +def test_audiox_model(model_name: str): + m = Omni(model=model_name, model_class_name="AudioXPipeline") + + # Keep runtime short for CI. + seconds_total = 2.0 + # AudioXPipeline always emits 44.1 kHz stereo (advertised via class-level + # ``audio_sample_rate``); the trimmed output should match this rate. + sample_rate = 44100 + + outputs = m.generate( + prompts={"prompt": "A dog barking in a quiet park."}, + sampling_params_list=OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=6.0, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), + num_outputs_per_prompt=1, + extra_args={ + "audiox_task": "t2a", + "seconds_start": 0.0, + "seconds_total": seconds_total, + }, + ), + ) + + assert outputs is not None + first_output = outputs[0] + assert first_output.final_output_type == "audio" + assert hasattr(first_output, "request_output") and first_output.request_output + + req_out = first_output.request_output + assert isinstance(req_out, OmniRequestOutput) + assert req_out.final_output_type == "audio" + assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output + + audio = req_out.multimodal_output.get("audio") + assert isinstance(audio, np.ndarray) + # audio shape: (batch, channels, samples) + assert audio.ndim == 3 + assert audio.shape[0] == 1 + assert audio.shape[1] == 2 + assert audio.shape[2] > 0 + expected_samples = int(seconds_total * sample_rate) + assert abs(audio.shape[2] - expected_samples) <= 2 * 1024 diff --git a/tests/e2e/offline_inference/test_ernie_image.py b/tests/e2e/offline_inference/test_ernie_image.py new file mode 100644 index 000000000..c790bbf7c --- /dev/null +++ b/tests/e2e/offline_inference/test_ernie_image.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for ErnieImage text-to-image generation. + +Equivalent to running: + vllm serve baidu/ERNIE-Image --omni +""" + +import pytest +from PIL import Image + +from tests.helpers.runtime import OmniRunner +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "baidu/ERNIE-Image" + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_ernie_image_text_to_image(): + with OmniRunner( + MODEL, + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + ), + enable_cpu_offload=True, + ) as runner: + omni_outputs = list( + runner.omni.generate( + prompts=["A photo of a cat sitting on a laptop"], + sampling_params_list=OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=2, + guidance_scale=4.0, + seed=42, + ), + ) + ) + + assert len(omni_outputs) > 0 + output = omni_outputs[0] + images = None + if output.images: + images = output.images + elif hasattr(output, "request_output") and output.request_output: + for stage_out in output.request_output: + if hasattr(stage_out, "images") and stage_out.images: + images = stage_out.images + break + + assert images is not None + assert len(images) > 0 + assert isinstance(images[0], Image.Image) + assert images[0].size == (512, 512) diff --git a/tests/e2e/offline_inference/test_internvla_a1.py b/tests/e2e/offline_inference/test_internvla_a1.py new file mode 100644 index 000000000..38104dc66 --- /dev/null +++ b/tests/e2e/offline_inference/test_internvla_a1.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Offline E2E test for InternVLA-A1 open-loop inference.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "internvla_a1" / "end2end.py" + + +def _required_env(name: str) -> str: + value = os.getenv(name) + if not value: + pytest.skip(f"{name} is not set") + return value + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) +def test_internvla_a1_offline_open_loop(run_level: str) -> None: + if run_level != "advanced_model": + pytest.skip("InternVLA-A1 offline evaluation requires real local checkpoints and dataset.") + + model_dir = _required_env("INTERNVLA_A1_MODEL_DIR") + dataset_dir = _required_env("INTERNVLA_A1_DATASET_DIR") + processor_dir = _required_env("INTERNVLA_A1_PROCESSOR_DIR") + cosmos_dir = _required_env("INTERNVLA_A1_COSMOS_DIR") + + env = os.environ.copy() + env["INTERNVLA_A1_PROCESSOR_DIR"] = processor_dir + env["INTERNVLA_A1_COSMOS_DIR"] = cosmos_dir + + with tempfile.TemporaryDirectory(prefix="internvla_a1_e2e_") as tmpdir: + output_dir = Path(tmpdir) / "outputs" + cmd = [ + sys.executable, + str(EXAMPLE_SCRIPT), + "--model-dir", + model_dir, + "--dataset-dir", + dataset_dir, + "--output-dir", + str(output_dir), + "--num-samples", + "1", + "--num-episodes", + "1", + "--dtype", + "float32", + "--attn-implementation", + "eager", + "--skip-plots", + ] + subprocess.run(cmd, check=True, env=env) + + summary_path = output_dir / "summary.json" + assert summary_path.exists(), f"Missing summary file: {summary_path}" + + with open(summary_path, encoding="utf-8") as f: + summary = json.load(f) + + assert summary["mode"] == "registry" + assert summary["attn_implementation"] == "eager" + assert summary["dtype"] == "float32" + assert len(summary["results"]) == 1 + assert "registry" in summary["eval_summaries"] + + eval_summary = summary["eval_summaries"]["registry"] + assert eval_summary["num_episodes"] == 1 + assert eval_summary["average_mse"] is not None + assert eval_summary["average_mae"] is not None + + log_path = output_dir / "registry" / "log.json" + assert log_path.exists(), f"Missing evaluation log: {log_path}" diff --git a/tests/e2e/offline_inference/test_moss_tts_nano.py b/tests/e2e/offline_inference/test_moss_tts_nano.py new file mode 100644 index 000000000..103f2834f --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_nano.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Nano single-stage pipeline. + +Every request needs a reference audio clip. We test upstream's +recommended ``voice_clone`` mode (no transcript needed). We fetch the +upstream sample (assets/audio/zh_1.wav, ~50 KB) once per test session +and reuse it for all cases. +""" + +from __future__ import annotations + +import os +import urllib.request + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from vllm_omni import Omni + +MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano" +SAMPLE_RATE = 48000 +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav" + +DEFAULT_SAMPLING = SamplingParams( + temperature=1.0, + top_p=1.0, + top_k=50, + max_tokens=4096, + seed=42, + detokenize=False, +) + + +@pytest.fixture(scope="session") +def ref_audio_path(tmp_path_factory) -> str: + """Download the upstream reference clip once per test session. + + Bounded with a 30 s timeout (urlretrieve has no timeout kwarg, so we use + urlopen + write ourselves). Fetch failures are hard failures, not skips, + so a broken network path can't silently mask regressions. Set + ``MOSS_TTS_NANO_SKIP_ON_NET_FAIL=1`` to opt into skipping for air-gapped + environments. + """ + cache_dir = tmp_path_factory.mktemp("moss_tts_nano_ref") + target = cache_dir / "zh_1.wav" + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + data = resp.read() + target.write_bytes(data) + except Exception as e: + msg = f"Cannot fetch upstream reference clip {REF_AUDIO_URL}: {e}" + if os.environ.get("MOSS_TTS_NANO_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not target.exists() or os.path.getsize(target) == 0: + pytest.fail(f"Reference clip empty after download: {target}") + return str(target) + + +def _build_request( + text: str, + prompt_audio_path: str, + prompt_text: str | None = None, + mode: str = "voice_clone", + max_new_frames: int = 100, # short for tests + seed: int = 42, +) -> dict: + """Build a MOSS-TTS-Nano offline request. + + Upstream forbids ``prompt_text`` in ``voice_clone`` mode; only forward + it when explicitly supplied (and typically with ``mode='continuation'``). + """ + additional: dict = { + "text": [text], + "mode": [mode], + "prompt_audio_path": [prompt_audio_path], + "max_new_frames": [max_new_frames], + "seed": [seed], + } + if prompt_text is not None: + additional["prompt_text"] = [prompt_text] + return { + "prompt": "<|im_start|>assistant\n", + "additional_information": additional, + } + + +def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: + """Run a single request and return (waveform, sample_rate).""" + for stage_outputs in omni.generate(request, DEFAULT_SAMPLING): + for req_output in stage_outputs.request_output: + mm = req_output.outputs[0].multimodal_output + assert mm is not None, "Expected multimodal_output to be non-None" + audio = mm.get("audio") + sr = mm.get("sr") + assert audio is not None, "Expected 'audio' key in multimodal_output" + assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}" + return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE + raise AssertionError("No stage outputs received") + + +@pytest.fixture(scope="module") +def omni_engine(): + """Module-scoped Omni engine to avoid re-loading the model for each test.""" + return Omni(model=MODEL_NAME, stage_init_timeout=180) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_english(omni_engine, ref_audio_path): + """English voice_clone mode (no prompt_text) produces non-empty 48 kHz audio.""" + req = _build_request("Hello, this is a short voice cloning demo for testing.", ref_audio_path) + audio, sr = _collect_audio(omni_engine, req) + + assert sr == SAMPLE_RATE, f"Expected sample_rate={SAMPLE_RATE}, got {sr}" + assert audio.numel() > 0, "Audio tensor should not be empty" + assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_chinese(omni_engine, ref_audio_path): + """Chinese voice_clone mode produces non-empty audio.""" + req = _build_request("你好，这是语音合成测试。", ref_audio_path) + audio, sr = _collect_audio(omni_engine, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_deterministic(omni_engine, ref_audio_path): + """Same seed produces identical waveforms.""" + req = _build_request("Reproducible output test.", ref_audio_path, seed=123) + audio1, _ = _collect_audio(omni_engine, req) + audio2, _ = _collect_audio(omni_engine, req) + + assert audio1.shape == audio2.shape, "Waveform shapes should match with same seed" + assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_nano_batch(omni_engine, ref_audio_path): + """Batch of two requests returns audio for each.""" + requests = [ + _build_request("First request.", ref_audio_path), + _build_request("Second request.", ref_audio_path), + ] + results = [] + for stage_outputs in omni_engine.generate(requests, [DEFAULT_SAMPLING] * 2): + for req_output in stage_outputs.request_output: + mm = req_output.outputs[0].multimodal_output + assert mm is not None + results.append(mm["audio"].cpu()) + + assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" + for i, audio in enumerate(results): + assert audio.numel() > 0, f"Audio {i} is empty" diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index 2763adf4f..0cc8deea2 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -32,6 +32,7 @@ _OMNI_RUNNER_PARAM = ( ) pytestmark = [ + pytest.mark.skip(reason="issue#4361"), pytest.mark.full_model, pytest.mark.tts, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), diff --git a/tests/e2e/online_serving/test_audiox_online.py b/tests/e2e/online_serving/test_audiox_online.py new file mode 100644 index 000000000..6ef84b206 --- /dev/null +++ b/tests/e2e/online_serving/test_audiox_online.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online serving test for AudioX text-to-audio diffusion. + +Mirrors `tests/e2e/online_serving/test_sd3_expansion.py` for image diffusion: spin up +`vllm-omni` with `--model-class-name AudioXPipeline` and validate via the standard +`send_diffusion_request` helper (now audio-aware). +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data + +# Tiny / random checkpoint usable in CI; override to a real bundle locally. +AUDIOX_TEST_MODEL = os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random") +T2A_PROMPT = "A quiet living room with soft fabric rustle and gentle cat breathing." + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) + + +def _audiox_server_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=["--model-class-name", "AudioXPipeline"], + ), + id="t2a", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True) +def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`.""" + request_config = { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT), + "extra_body": { + "num_inference_steps": 4, + "guidance_scale": 6.0, + "seed": 42, + "audiox_task": "t2a", + "seconds_start": 0.0, + "seconds_total": 2.0, + }, + } + openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py index 8ab327022..add53e3d9 100644 --- a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py +++ b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py @@ -26,6 +26,12 @@ _SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip( reason="Requires 4x H100 GPUs; skipped in CI for now.", ) +pytestmark = [pytest.mark.omni, pytest.mark.full_model] + +_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip( + reason="Requires 4x H100 GPUs; skipped in CI for now.", +) + models = ["Jonathan1909/Ming-flash-omni-2.0"] # Use thinker-only topology to test text outputs diff --git a/tests/e2e/online_serving/test_moss_tts_nano.py b/tests/e2e/online_serving/test_moss_tts_nano.py new file mode 100644 index 000000000..e89571988 --- /dev/null +++ b/tests/e2e/online_serving/test_moss_tts_nano.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for MOSS-TTS-Nano model with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. The server uses upstream's +``voice_clone`` mode (no transcript needed); ``ref_text`` and ``voice`` +are accepted in the OpenAI schema but ignored. One case includes a +non-empty ``ref_text`` to verify it does not break the pipeline. The +reference clip is fetched from the upstream repo (assets/audio/zh_1.wav, +~50 KB). +""" + +import base64 +import os +import urllib.request + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "OpenMOSS-Team/MOSS-TTS-Nano" +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav" +REF_AUDIO_TRANSCRIPT = "欢迎关注模思智能、上海创智学院与复旦大学自然语言处理实验室。" + + +@pytest.fixture(scope="session") +def ref_audio_data_url() -> str: + """Fetch the upstream sample clip and return it as a base64 data URL. + + The fetch failure is escalated to a hard failure (not pytest.skip) so that + a broken network path does not silently mask regressions in + /v1/audio/speech. Set ``MOSS_TTS_NANO_SKIP_ON_NET_FAIL=1`` to opt into + skipping in air-gapped environments. + """ + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + data = resp.read() + except Exception as e: + msg = f"Cannot fetch upstream reference clip {REF_AUDIO_URL}: {e}" + if os.environ.get("MOSS_TTS_NANO_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not data: + pytest.fail(f"Reference clip empty: {REF_AUDIO_URL}") + return f"data:audio/wav;base64,{base64.b64encode(data).decode('ascii')}" + + +def get_prompt(prompt_type="text"): + """Text prompt for text-to-audio tests. + + Avoid the model's own name ("MOSS-TTS-Nano") in the test input — the + codec consistently mispronounces it ("MOS's TTS Nano", "Was TTS Nano", + etc.), which trips the transcript-similarity assertion in CI without + indicating a real regression. Use plain natural-sounding sentences. + """ + prompts = { + "text": "Hello, this is a short voice cloning demo for testing.", + "chinese": "你好，这是一段简单的语音合成测试。", + } + return prompts.get(prompt_type, prompts["text"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("moss_tts_nano.yaml"), + server_args=["--disable-log-stats"], + ), + id="moss_tts_nano", + ) +] + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None: + """ + Test voice_clone mode (no ref_text) via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text + reference audio (no transcript) + Output Modal: audio (48 kHz, WAV) + Input Setting: stream=False + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "response_format": "wav", + "ref_audio": ref_audio_data_url, + } + + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> None: + """ + Test streaming voice_clone via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text + reference audio (no transcript) + Output Modal: audio (48 kHz, PCM stream) + Input Setting: stream=True + Datasets: single request + + NOTE: ``min_hnr_db=-5.0`` overrides the conftest default of 1.0 dB. + MOSS-TTS-Nano's voice_clone output is intrinsically noisy by that + metric — verified on H20, both streaming and non-streaming clips + measure HNR around -2 dB even when the audio sounds correct. We + keep the PCM-stream check in place to catch catastrophic decoder + failures (which produce HNR much further below 0), just with a + threshold MOSS can clear in normal operation. + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "response_format": "pcm", + "ref_audio": ref_audio_data_url, + "min_hnr_db": -5.0, + } + + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> None: + """ + Test Chinese voice_clone via /v1/audio/speech. + Deploy Setting: default yaml + Input Modal: text (Chinese) + reference audio (no transcript) + Output Modal: audio (48 kHz, WAV) + Input Setting: stream=False + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("chinese"), + "stream": False, + "response_format": "wav", + "ref_audio": ref_audio_data_url, + } + + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audio_data_url) -> None: + """ + Sending ``ref_text`` alongside ``ref_audio`` is accepted but ignored. + + The OpenAI ``ref_text`` field is part of the schema, so clients may + send it. The server uses upstream's voice_clone mode regardless and + drops the transcript; this case verifies the request still completes + and produces audio. + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("chinese"), + "stream": False, + "response_format": "wav", + "ref_audio": ref_audio_data_url, + "ref_text": REF_AUDIO_TRANSCRIPT, + } + + openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py b/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py new file mode 100644 index 000000000..f3b261081 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out). +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +import os +import wave + +import pytest +import websockets + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + convert_audio_bytes_to_text, + cosine_similarity_text, + generate_synthetic_audio, +) +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + +# The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU +# via its ``platforms:`` section, so one path serves all three. +default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +def _realtime_stage_config_path() -> str: + """CI omni layout without async_chunk; stage 0 thinker max_tokens=10.""" + return modify_stage_config( + default_stage_config, + updates={"stages": {0: {"default_sampling_params.max_tokens": 10}}}, + ) + + +realtime_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=_realtime_stage_config_path(), + use_stage_cli=True, + ), + id="thinker_max_tokens_10", + ), +] + + +def _pcm16_mono_16k_from_wav_bytes(wav_bytes: bytes) -> bytes: + with wave.open(io.BytesIO(wav_bytes), "rb") as wf: + if wf.getnchannels() != 1: + raise ValueError(f"Expected mono WAV, got {wf.getnchannels()} channels") + if wf.getsampwidth() != 2: + raise ValueError(f"Expected 16-bit PCM, sampwidth={wf.getsampwidth()}") + if wf.getframerate() != 16000: + raise ValueError(f"Expected 16 kHz input for /v1/realtime, got {wf.getframerate()} Hz") + if wf.getcomptype() != "NONE": + raise ValueError(f"Expected uncompressed PCM, comptype={wf.getcomptype()!r}") + return wf.readframes(wf.getnframes()) + + +def _wav_bytes_from_pcm16(pcm: bytes, sample_rate_hz: int) -> bytes: + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate_hz) + wf.writeframes(pcm) + return buf.getvalue() + + +async def _run_realtime_audio_roundtrip( + host: str, + port: int, + model: str, + pcm16: bytes, + *, + chunk_ms: int = 100, +) -> dict: + uri = f"ws://{host}:{port}/v1/realtime" + incremental: list[bytes] = [] + output_sr = 24000 + text_chunks: list[str] = [] + final_text = "" + delta_events = 0 + + bytes_per_ms = 16000 * 2 // 1000 + chunk_bytes = max(bytes_per_ms * chunk_ms, 2) + + async with websockets.connect(uri, max_size=64 * 1024 * 1024) as ws: + await ws.send(json.dumps({"type": "session.update", "model": model})) + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) + + for i in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[i : i + chunk_bytes] + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("utf-8"), + } + ) + ) + + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + + while True: + message = await asyncio.wait_for(ws.recv(), timeout=600) + if isinstance(message, bytes): + continue + + event = json.loads(message) + event_type = event.get("type") + + if event_type == "session.created": + continue + + if event_type == "response.audio.delta": + delta_events += 1 + sr = event.get("sample_rate_hz") + if isinstance(sr, int) and sr > 0: + output_sr = sr + audio_b64 = event.get("audio", "") + if audio_b64: + incremental.append(base64.b64decode(audio_b64)) + continue + + if event_type == "transcription.delta": + d = event.get("delta", "") + if d: + text_chunks.append(d) + continue + + if event_type == "transcription.done": + final_text = event.get("text", "") or "".join(text_chunks) + continue + + if event_type == "response.audio.done": + break + + if event_type == "error": + raise AssertionError(f"WebSocket error: {event}") + + raise AssertionError(f"Unexpected WebSocket event: {event}") + + out_pcm = b"".join(incremental) + return { + "output_pcm": out_pcm, + "output_sample_rate": output_sr, + "transcription_text": final_text if final_text else "".join(text_chunks), + "delta_events": delta_events, + } + + +class TestQwen3OmniRealtimeWebSocket: + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + @pytest.mark.parametrize("omni_server", realtime_server_params, indirect=True) + def test_streaming_audio_input_pcm_output(self, omni_server) -> None: + """ + Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and + transcription. Verify Whisper(output audio) aligns with model text (same idea + as multimodal omni e2e). + """ + syn = generate_synthetic_audio(10, 1, sample_rate=16000) + wav_bytes = base64.b64decode(syn["base64"]) + pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes) + + result = asyncio.run( + _run_realtime_audio_roundtrip( + omni_server.host, + omni_server.port, + omni_server.model, + pcm16, + chunk_ms=100, + ) + ) + + out_pcm = result["output_pcm"] + assert result["delta_events"] >= 1 + assert out_pcm, "No output PCM from response.audio.delta" + assert len(out_pcm) % 2 == 0 + assert len(out_pcm) >= 4096, "Output audio unexpectedly small" + assert result["output_sample_rate"] > 0 + + final_text = (result["transcription_text"] or "").strip() + assert final_text, "Expected non-empty transcription (model text stream)" + + wav_out = _wav_bytes_from_pcm16(out_pcm, result["output_sample_rate"]) + whisper_text = convert_audio_bytes_to_text(wav_out).strip() + assert whisper_text, "Whisper returned empty string for synthesized output audio" + + sim = cosine_similarity_text(whisper_text.lower(), final_text.lower()) + assert sim > 0.9, ( + f"Output audio transcript should match model text (sim={sim:.3f}): " + f"whisper={whisper_text!r}, model_text={final_text!r}" + ) diff --git a/tests/examples/offline_inference/test_qwen3_tts_mr_v2.py b/tests/examples/offline_inference/test_qwen3_tts_mr_v2.py new file mode 100644 index 000000000..89ccc094e --- /dev/null +++ b/tests/examples/offline_inference/test_qwen3_tts_mr_v2.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Offline inference tests: Qwen3-TTS. +See examples/offline_inference/qwen3_tts/README.md + +Picks classic test examples covering all three query types (CustomVoice, +VoiceDesign, Base) and the streaming execution path. +""" + +import tempfile +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from tests.examples.helpers import EXAMPLES, run_cmd +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.advanced_model, pytest.mark.example] + +END2END = str(EXAMPLES / "offline_inference" / "text_to_speech" / "qwen3_tts" / "end2end.py") + +CUSTOM_VOICE_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" +VOICE_DESIGN_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" +BASE_MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" + + +MIN_DURATION_S = 0.5 +EXPECTED_SAMPLE_RATE = 24000 +MR_V2_ENV = { + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "VLLM_OMNI_USE_V2_RUNNER": "1", +} + + +def _assert_wav_output(output_dir: str) -> None: + """Assert that at least one valid .wav file was produced. + + Checks: file exists, sample rate, minimum duration, and not silence. + """ + wav_files = list(Path(output_dir).glob("*.wav")) + assert len(wav_files) > 0, f"No .wav files found in {output_dir}" + for wav in wav_files: + data, sr = sf.read(wav, dtype="float32") + assert sr == EXPECTED_SAMPLE_RATE, f"{wav.name}: sample_rate={sr}, expected {EXPECTED_SAMPLE_RATE}" + duration = len(data) / sr + assert duration >= MIN_DURATION_S, f"{wav.name}: duration={duration:.3f}s, expected >= {MIN_DURATION_S}s" + assert np.max(np.abs(data)) > 0.01, f"{wav.name}: audio appears to be silence" + + +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_custom_voice(): + """CustomVoice single prompt — the most common TTS use case.""" + with tempfile.TemporaryDirectory() as output_dir: + command = [ + "python", + END2END, + "--model", + CUSTOM_VOICE_MODEL, + "--query-type", + "CustomVoice", + "--output-dir", + output_dir, + ] + run_cmd(command, env=MR_V2_ENV) + _assert_wav_output(output_dir) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_voice_design(): + """VoiceDesign single prompt — generates speech from a voice description.""" + with tempfile.TemporaryDirectory() as output_dir: + command = [ + "python", + END2END, + "--model", + VOICE_DESIGN_MODEL, + "--query-type", + "VoiceDesign", + "--output-dir", + output_dir, + ] + run_cmd(command, env=MR_V2_ENV) + _assert_wav_output(output_dir) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_base_icl(): + """Base ICL mode — voice cloning with reference audio and transcript.""" + with tempfile.TemporaryDirectory() as output_dir: + command = [ + "python", + END2END, + "--model", + BASE_MODEL, + "--query-type", + "Base", + "--mode-tag", + "icl", + "--output-dir", + output_dir, + ] + run_cmd(command, env=MR_V2_ENV) + _assert_wav_output(output_dir) + + +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_custom_voice_streaming(): + """CustomVoice streaming — exercises the AsyncOmni streaming path.""" + with tempfile.TemporaryDirectory() as output_dir: + command = [ + "python", + END2END, + "--model", + CUSTOM_VOICE_MODEL, + "--query-type", + "CustomVoice", + "--streaming", + "--output-dir", + output_dir, + ] + run_cmd(command, env=MR_V2_ENV) + _assert_wav_output(output_dir) diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index 0d0ec6f46..5c82e11d7 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -115,13 +115,16 @@ def test_send_multimodal_request_002(omni_server) -> None: @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: + env = os.environ.copy() + env["MODEL"] = omni_server.model + command = [ "bash", os.path.join(example_dir, "qwen2_5_omni/run_curl_multimodal_generation.sh"), "mixed_modalities", ] - result = run_cmd(command) + result = run_cmd(command, env=env) text_content = extract_content_after_keyword("Output of request:", result) diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index da62ceab2..6f9ded0b2 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -541,6 +541,15 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], assert response.success, "The request failed." + # Optional floor on decoded audio size (models with very short clips may use a lower value). + min_audio = request_config.get("min_audio_bytes") + if min_audio is not None: + n = int(min_audio) + if n > 0: + ab = response.audio_bytes + assert ab is not None, "Expected audio bytes when min_audio_bytes is set" + assert len(ab) > n, f"Audio payload too small: {len(ab)} bytes, expected more than {n} (min_audio_bytes)" + # Optional floor on decoded audio size (models with very short clips may use a lower value). min_audio = request_config.get("min_audio_bytes") if min_audio is not None: diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_conditioning.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_conditioning.py new file mode 100644 index 000000000..76d0651a4 --- /dev/null +++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_talker_conditioning.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import ( + Qwen2_5OmniForConditionalGeneration, +) + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + pytest.mark.skip( + reason="Methods _build_thinker_to_talker_latent and _build_talker_decode_reply_cache not yet implemented in Qwen2_5OmniForConditionalGeneration" + ), +] + + +def test_build_thinker_to_talker_latent_adds_token_embeddings(): + model = object.__new__(Qwen2_5OmniForConditionalGeneration) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + embeds = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + + latent = model._build_thinker_to_talker_latent(hidden, embeds, input_ids=None) + + assert torch.equal(latent, hidden + embeds) + + +def test_build_thinker_to_talker_latent_masks_multimodal_token_embeddings(): + model = object.__new__(Qwen2_5OmniForConditionalGeneration) + nn.Module.__init__(model) + model.thinker_config = SimpleNamespace( + audio_token_index=7, + image_token_index=8, + video_token_index=9, + ) + + class _FakeThinker: + @staticmethod + def embed_input_ids(input_ids: torch.Tensor) -> torch.Tensor: + values = input_ids.to(torch.float32) + return torch.stack([values, values + 100.0], dim=-1) + + model.thinker = _FakeThinker() + hidden = torch.zeros((4, 2), dtype=torch.float32) + input_ids = torch.tensor([1, 7, 2, 8], dtype=torch.long) + multimodal_embeds = torch.full((4, 2), 999.0, dtype=torch.float32) + + latent = model._build_thinker_to_talker_latent(hidden, multimodal_embeds, input_ids) + + assert torch.equal( + latent, + torch.tensor( + [ + [1.0, 101.0], + [0.0, 0.0], + [2.0, 102.0], + [0.0, 0.0], + ], + ), + ) + + +def test_build_talker_decode_reply_cache_appends_text_eos_pad(): + model = object.__new__(Qwen2_5OmniForConditionalGeneration) + model.embed_text_eos_token = torch.tensor([[100.0, 101.0]]) + model.embed_text_pad_token = torch.tensor([[200.0, 201.0]]) + thinker_result = torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]]) + + reply = model._build_talker_decode_reply_cache( + thinker_result, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + assert torch.equal( + reply, + torch.tensor( + [ + [2.0, 2.0], + [3.0, 3.0], + [100.0, 101.0], + [200.0, 201.0], + ] + ), + ) + + +def test_talker_decode_keeps_final_reply_pad_vector(): + model = object.__new__(Qwen2_5OmniForConditionalGeneration) + nn.Module.__init__(model) + model.model = nn.Linear(2, 2) + + class _FakeTalker: + @staticmethod + def embed_input_ids(input_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros((input_ids.numel(), 2), dtype=torch.float32) + + model.talker = _FakeTalker() + input_ids = torch.tensor([123], dtype=torch.long) + input_embeds = torch.zeros((1, 2), dtype=torch.float32) + final_pad = torch.tensor([[5.0, 6.0]], dtype=torch.float32) + + _, out_embeds, update = model.thinker_to_talker_decode_one_step( + input_ids, + input_embeds, + {"embed": {"thinker_reply": final_pad}}, + ) + + torch.testing.assert_close(out_embeds, final_pad) + assert update == {} + + +class _FakeToken2Wav(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(())) + self.calls = [] + self.token2wav = SimpleNamespace( + factor=SimpleNamespace(factor=2), + code2wav_dit_model=SimpleNamespace(mel_dim=3), + ) + + def forward( + self, + code: torch.Tensor, + *, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + num_steps: int = 10, + **_: object, + ) -> torch.Tensor: + self.calls.append( + { + "code": code.clone(), + "conditioning": conditioning.clone(), + "reference_mel": reference_mel.clone(), + "num_steps": num_steps, + } + ) + return torch.arange(6, dtype=torch.float32) + + def process_chunk(self, **_: object) -> tuple[torch.Tensor, torch.Tensor]: + raise AssertionError("_codec_to_audio should not truncate codec through process_chunk") + + +def test_codec_to_audio_uses_full_token2wav_forward(): + model = object.__new__(Qwen2_5OmniForConditionalGeneration) + nn.Module.__init__(model) + token2wav = _FakeToken2Wav() + cond = torch.ones((1, 4), dtype=torch.float32) + ref_mel = torch.full((1, 300, 3), 2.0, dtype=torch.float32) + model.token2wav = token2wav + model.token2wav_config = SimpleNamespace( + dit_config=SimpleNamespace(enc_emb_dim=4, mel_dim=3), + ) + model._token2wav_conds = {"Chelsie": cond} + model._token2wav_ref_mels = {"Chelsie": ref_mel} + + out = model._codec_to_audio(torch.arange(80, dtype=torch.long), "Chelsie") + + assert len(token2wav.calls) == 1 + call = token2wav.calls[0] + torch.testing.assert_close(call["code"], torch.arange(80, dtype=torch.long).unsqueeze(0)) + torch.testing.assert_close(call["conditioning"], cond) + torch.testing.assert_close(call["reference_mel"], ref_mel) + assert call["num_steps"] == 10 + torch.testing.assert_close(out, torch.arange(6, dtype=torch.float32)) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrope.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrope.py new file mode 100644 index 000000000..a344d046e --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrope.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( + Qwen3OmniMoeForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_non_thinker_mrope_positions_do_not_require_multimodal_kwargs(): + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + model.model_stage = "talker" + + positions, delta = model.get_mrope_input_positions([1, 2, 3]) + + assert delta == 0 + assert torch.equal( + positions, + torch.tensor( + [ + [0, 1, 2], + [0, 1, 2], + [0, 1, 2], + ] + ), + ) diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni.py new file mode 100644 index 000000000..4a64ad699 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import ( + TALKER_CODEC_END_TOKEN_ID, + TALKER_CODEC_PAD_TOKEN_ID, + TALKER_CODEC_START_TOKEN_ID, + talker2code2wav, +) + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + pytest.mark.skip( + reason="qwen2_5_omni stage input processor tests need alignment with model changes that are not yet implemented" + ), +] + + +def _source_outputs(output): + return [SimpleNamespace(outputs=[output])] + + +def test_talker2code2wav_accepts_completion_output_token_ids(): + output = SimpleNamespace(token_ids=[TALKER_CODEC_START_TOKEN_ID, 11, 22, TALKER_CODEC_END_TOKEN_ID]) + + prompts = talker2code2wav(_source_outputs(output)) + + assert [prompt["prompt_token_ids"] for prompt in prompts] == [[11, 22]] + + +def test_talker2code2wav_prefers_cumulative_token_ids_when_present(): + output = SimpleNamespace( + token_ids=[999], + cumulative_token_ids=[TALKER_CODEC_START_TOKEN_ID, 33, 44, TALKER_CODEC_END_TOKEN_ID], + ) + + prompts = talker2code2wav(_source_outputs(output)) + + assert [prompt["prompt_token_ids"] for prompt in prompts] == [[33, 44]] + + +def test_talker2code2wav_strips_terminal_codec_pad(): + output = SimpleNamespace(cumulative_token_ids=[TALKER_CODEC_START_TOKEN_ID, 55, 66, TALKER_CODEC_PAD_TOKEN_ID]) + + prompts = talker2code2wav(_source_outputs(output)) + + assert [prompt["prompt_token_ids"] for prompt in prompts] == [[55, 66]] + + +def test_talker2code2wav_keeps_empty_codec_terminal_input(): + output = SimpleNamespace(cumulative_token_ids=[TALKER_CODEC_START_TOKEN_ID, TALKER_CODEC_END_TOKEN_ID]) + + prompts = talker2code2wav(_source_outputs(output)) + + assert [prompt["prompt_token_ids"] for prompt in prompts] == [[]] diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 52ca29905..d93b1801c 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections import defaultdict from types import SimpleNamespace import pytest @@ -76,6 +77,26 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex assert state.talker2code2wav_last_seq_len == 0 +def test_talker2code2wav_async_chunk_emits_eof_marker_when_finished_empty() -> None: + transfer_manager = SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}), + ) + request = SimpleNamespace(external_req_id="r1", is_finished=lambda: True) + + for multimodal_output in (None, {"codes": {"audio": torch.empty((0,), dtype=torch.long)}}): + payload = q3.talker2code2wav_async_chunk( + transfer_manager=transfer_manager, + multimodal_output=multimodal_output, + request=request, + is_finished=True, + ) + + assert payload["codes"] == {"audio": []} + assert payload["meta"]["left_context_size"] == 0 + assert payload["meta"]["finished"].item() is True + + def test_streaming_input_prefill_chunk_is_cached() -> None: transfer_manager = SimpleNamespace(_pending_streaming_prefills={}) request = SimpleNamespace( diff --git a/tests/worker/test_gpu_ar_model_runner_bookkeeping_compat.py b/tests/worker/test_gpu_ar_model_runner_bookkeeping_compat.py new file mode 100644 index 000000000..566ac2f69 --- /dev/null +++ b/tests/worker/test_gpu_ar_model_runner_bookkeeping_compat.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + +def test_gpu_ar_model_runner_calls_bookkeeping_with_current_vllm_signature(): + runner = object.__new__(GPUARModelRunner) + calls = [] + + def bookkeeping_sync(*args): + calls.append(args) + return (0, None, [], {}, [], {}, []) + + runner._bookkeeping_sync = bookkeeping_sync + runner._bookkeeping_accepts_spec_decode_metadata = False + scheduler_output = SimpleNamespace(total_num_scheduled_tokens=3) + sampler_output = SimpleNamespace() + logits = torch.empty(1, 1) + hidden_states = torch.empty(1, 1) + + result = runner._bookkeeping_sync_compat( + scheduler_output, + sampler_output, + logits, + hidden_states, + scheduler_output.total_num_scheduled_tokens, + spec_decode_metadata=object(), + ) + + assert result == (0, None, [], {}, [], {}, []) + assert calls == [ + ( + scheduler_output, + sampler_output, + logits, + hidden_states, + 3, + ) + ] + + +if __name__ == "__main__": + test_gpu_ar_model_runner_calls_bookkeeping_with_current_vllm_signature() + + +def test_bookkeeping_compat_does_not_reflect_per_call(): + import inspect + + source = inspect.getsource(GPUARModelRunner._bookkeeping_sync_compat) + assert "inspect.signature" not in source + + +def test_sample_tokens_passes_spec_decode_metadata_to_bookkeeping_compat(): + import ast + import inspect + + source = inspect.getsource(GPUARModelRunner.sample_tokens) + tree = ast.parse(source) + + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_bookkeeping_sync_compat" + ] + assert len(calls) == 1 + assert any(isinstance(arg, ast.Name) and arg.id == "spec_decode_metadata" for arg in calls[0].args) diff --git a/tests/worker/test_omni_generation_scheduler_vllm20_compat.py b/tests/worker/test_omni_generation_scheduler_vllm20_compat.py new file mode 100644 index 000000000..4b3c42f22 --- /dev/null +++ b/tests/worker/test_omni_generation_scheduler_vllm20_compat.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +from vllm_omni.core.sched.omni_generation_scheduler import ( + _get_request_num_cached_tokens, + _make_engine_core_output, +) + + +def test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing(): + request = SimpleNamespace(request_id="rid", num_computed_tokens=5) + + assert _get_request_num_cached_tokens(request) == 0 + + +def test_make_engine_core_output_ignores_fields_missing_from_current_vllm(): + output = _make_engine_core_output( + request_id="rid", + new_token_ids=[1], + finish_reason=None, + new_logprobs=None, + new_prompt_logprobs_tensors=None, + pooling_output=None, + stop_reason=None, + events=[], + kv_transfer_params=None, + trace_headers={}, + num_cached_tokens=3, + num_external_computed_tokens=4, + routed_experts=None, + num_nans_in_logits=0, + ) + + assert output.request_id == "rid" + assert output.new_token_ids == [1] + + +if __name__ == "__main__": + test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing() + test_make_engine_core_output_ignores_fields_missing_from_current_vllm() diff --git a/tests/worker/test_omni_generation_worker_compile_time.py b/tests/worker/test_omni_generation_worker_compile_time.py new file mode 100644 index 000000000..875b17943 --- /dev/null +++ b/tests/worker/test_omni_generation_worker_compile_time.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect + +import vllm_omni.worker.gpu_generation_worker as worker +from vllm_omni.worker.gpu_generation_worker import _make_compilation_times + + +def test_make_compilation_times_matches_current_vllm_shape(): + result = _make_compilation_times(0.0) + + if isinstance(result, float): + assert result == 0.0 + else: + assert result.language_model == 0.0 + assert result.encoder == 0.0 + + +if __name__ == "__main__": + test_make_compilation_times_matches_current_vllm_shape() + + +def test_make_compilation_times_filters_to_current_fields(): + result = _make_compilation_times(1.5, speculative_model=2.0) + + if isinstance(result, float): + assert result == 1.5 + else: + assert result.language_model == 1.5 + assert result.encoder == 0.0 + + +def test_compilation_times_import_is_lazy(): + source_before_helper = inspect.getsource(worker).split("def _make_compilation_times", 1)[0] + assert "CompilationTimes" not in source_before_helper diff --git a/tests/worker/test_omni_scheduler_engine_core_output_compat.py b/tests/worker/test_omni_scheduler_engine_core_output_compat.py new file mode 100644 index 000000000..ee586d570 --- /dev/null +++ b/tests/worker/test_omni_scheduler_engine_core_output_compat.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +from vllm_omni.core.sched.omni_ar_scheduler import ( + _get_request_num_cached_tokens, + _get_request_num_external_computed_tokens, + _make_engine_core_output, +) + + +def test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing(): + request = SimpleNamespace(request_id="rid", num_computed_tokens=5) + + assert _get_request_num_cached_tokens(request) == 0 + + +def test_get_request_num_external_computed_tokens_defaults_to_zero_when_field_is_missing(): + request = SimpleNamespace(request_id="rid") + + assert _get_request_num_external_computed_tokens(request) == 0 + + +def test_make_engine_core_output_ignores_fields_missing_from_current_vllm(): + output = _make_engine_core_output( + request_id="rid", + new_token_ids=[1], + finish_reason=None, + new_logprobs=None, + new_prompt_logprobs_tensors=None, + pooling_output=None, + stop_reason=None, + events=[], + kv_transfer_params=None, + trace_headers={}, + num_cached_tokens=3, + num_external_computed_tokens=4, + routed_experts=None, + num_nans_in_logits=0, + is_segment_finished=None, + new_prompt_len_snapshot=None, + ) + + assert output.request_id == "rid" + assert output.new_token_ids == [1] + + +if __name__ == "__main__": + test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing() + test_get_request_num_external_computed_tokens_defaults_to_zero_when_field_is_missing() + test_make_engine_core_output_ignores_fields_missing_from_current_vllm() diff --git a/tests/worker/test_omni_v2_async_chunk_updates.py b/tests/worker/test_omni_v2_async_chunk_updates.py new file mode 100644 index 000000000..bd8a789fb --- /dev/null +++ b/tests/worker/test_omni_v2_async_chunk_updates.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm_omni.core.sched.output import OmniCachedRequestData +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + +class _Array: + def __init__(self): + self.np = [0] + self.writes = [] + self.applied = False + + def stage_write_elem(self, idx, value): + self.writes.append((idx, value)) + + def apply_write(self): + self.applied = True + + +class _TokenIds: + def __init__(self): + self.writes = [] + + def stage_write(self, idx, start, values): + self.writes.append((idx, start, list(values))) + + +class _ReqStates: + def __init__(self, has_request=True): + self.req_id_to_index = {"r1": 0} if has_request else {} + self.prompt_len = SimpleNamespace(np=[0]) + self.prefill_len = SimpleNamespace(np=[0]) + self.total_len = _Array() + self.all_token_ids = _TokenIds() + self.num_computed_tokens = _Array() + self.num_computed_prefill_tokens = [7] + self.applied = False + self.removed = [] + + def apply_staged_writes(self): + self.applied = True + + def remove_request(self, req_id): + self.removed.append(req_id) + self.req_id_to_index.pop(req_id, None) + + +class _IntermediateBuffer: + def __init__(self): + self.buffers = [{"stale": "value", "req_id": "r1"}] + + def remove_request(self, idx): + self.buffers[idx] = {} + + +def test_async_chunk_update_clears_stale_intermediate_buffer(): + runner = object.__new__(OmniGenerationModelRunner) + runner.req_states = _ReqStates() + runner.model_state = SimpleNamespace(intermediate_buffer=_IntermediateBuffer()) + cached = OmniCachedRequestData( + req_ids=["r1"], + resumed_req_ids=set(), + new_token_ids=[], + all_token_ids=[], + new_block_ids=[], + num_computed_tokens=[], + num_output_tokens=[], + prompt_token_ids={"r1": [1, 2, 3]}, + additional_information={"r1": {"fresh": "value"}}, + ) + + runner._handle_async_chunk_updates(SimpleNamespace(scheduled_cached_reqs=cached)) + + assert runner.model_state.intermediate_buffer.buffers[0] == {} + assert runner.req_states.prompt_len.np[0] == 3 + assert runner.req_states.prefill_len.np[0] == 3 + assert runner.req_states.all_token_ids.writes == [(0, 0, [1, 2, 3])] + assert runner.req_states.num_computed_prefill_tokens[0] == 0 + assert runner.req_states.applied is True + + +def test_async_chunk_update_readds_released_cached_request(): + runner = object.__new__(OmniGenerationModelRunner) + runner.req_states = _ReqStates(has_request=False) + runner.model_state = SimpleNamespace(intermediate_buffer=_IntermediateBuffer()) + added = [] + runner.add_requests = lambda scheduler_output: added.extend(scheduler_output.scheduled_new_reqs) + cached = OmniCachedRequestData( + req_ids=["r1"], + resumed_req_ids=set(), + new_token_ids=[], + all_token_ids={"r1": [1, 2, 3]}, + new_block_ids=[([7],)], + num_computed_tokens=[0], + num_output_tokens=[0], + prompt_token_ids={"r1": [1, 2, 3]}, + additional_information={"r1": {"fresh": "value"}}, + ) + + runner._handle_async_chunk_updates(SimpleNamespace(scheduled_cached_reqs=cached)) + + assert len(added) == 1 + assert added[0].req_id == "r1" + assert added[0].prompt_token_ids == [1, 2, 3] + assert added[0].prefill_token_ids == [1, 2, 3] + assert added[0].block_ids == ([7],) + assert added[0].additional_information == {"fresh": "value"} + + +def _make_runner_for_sample(async_chunk): + runner = object.__new__(OmniGenerationModelRunner) + runner._gen_model_output = OmniOutput( + text_hidden_states=torch.zeros(1, 1), + multimodal_outputs={"audio": [torch.zeros(1)]}, + ) + runner._gen_input_batch = SimpleNamespace( + num_reqs=1, + idx_mapping_np=[0], + req_ids=["r1"], + ) + runner._gen_kv_connector_output = None + runner.execute_model_state = object() + runner.model_config = SimpleNamespace(async_chunk=async_chunk) + runner.req_states = _ReqStates() + runner.model_state = SimpleNamespace( + intermediate_buffer=_IntermediateBuffer(), + remove_request=lambda idx: runner.model_state.intermediate_buffer.remove_request(idx), + ) + removed = [] + runner._remove_request = lambda req_id: removed.append(req_id) or True + return runner, removed + + +def test_generation_sample_keeps_runner_slot_without_async_chunk(): + runner, removed = _make_runner_for_sample(async_chunk=False) + + output = runner.sample_tokens() + + assert output.req_ids == ["r1"] + assert removed == [] + assert runner.model_state.intermediate_buffer.buffers[0] == {"stale": "value", "req_id": "r1"} + + +def test_generation_sample_releases_runner_slot_after_chunk_output(): + runner, removed = _make_runner_for_sample(async_chunk=True) + + output = runner.sample_tokens() + + assert output.req_ids == ["r1"] + assert removed == ["r1"] + assert runner.model_state.intermediate_buffer.buffers[0] == {} + + +if __name__ == "__main__": + test_async_chunk_update_clears_stale_intermediate_buffer() + test_async_chunk_update_readds_released_cached_request() + test_generation_sample_keeps_runner_slot_without_async_chunk() + test_generation_sample_releases_runner_slot_after_chunk_output() diff --git a/tests/worker/test_omni_v2_async_output_signature.py b/tests/worker/test_omni_v2_async_output_signature.py new file mode 100644 index 000000000..6417b5435 --- /dev/null +++ b/tests/worker/test_omni_v2_async_output_signature.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect + +from vllm_omni.worker_v2.omni_ar_model_runner import OmniAsyncOutput + + +def test_omni_async_output_does_not_require_external_copy_event(): + param = inspect.signature(OmniAsyncOutput.__init__).parameters["copy_event"] + + assert param.default is None + + +if __name__ == "__main__": + test_omni_async_output_does_not_require_external_copy_event() diff --git a/tests/worker/test_omni_v2_execute_model_signature.py b/tests/worker/test_omni_v2_execute_model_signature.py new file mode 100644 index 000000000..3d06645cc --- /dev/null +++ b/tests/worker/test_omni_v2_execute_model_signature.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect + +from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner +from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + + +def test_v2_model_runners_accept_vllm_020_is_profile_keyword(): + for runner_cls in (OmniARModelRunner, OmniGenerationModelRunner, OmniGPUModelRunner): + assert "is_profile" in inspect.signature(runner_cls.execute_model).parameters + + +if __name__ == "__main__": + test_v2_model_runners_accept_vllm_020_is_profile_keyword() diff --git a/tests/worker/test_omni_v2_execute_model_state_compat.py b/tests/worker/test_omni_v2_execute_model_state_compat.py new file mode 100644 index 000000000..dff4f20b1 --- /dev/null +++ b/tests/worker/test_omni_v2_execute_model_state_compat.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.worker_v2.omni_model_runner import _make_execute_model_state + + +def test_make_execute_model_state_ignores_fields_missing_from_current_vllm(): + state = _make_execute_model_state( + input_batch="input", + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states="hidden", + aux_hidden_states=None, + # finished_req_ids is a real field on vLLM 0.23.0's ExecuteModelState. + finished_req_ids=set(), + # These are fields the omni runners still pass but current vLLM no longer + # declares; they must be filtered out rather than crash the constructor. + kv_connector_output=None, + num_tokens_across_dp=8, + ) + + assert state.input_batch == "input" + assert state.hidden_states == "hidden" + assert not hasattr(state, "num_tokens_across_dp") + assert not hasattr(state, "kv_connector_output") + + +if __name__ == "__main__": + test_make_execute_model_state_ignores_fields_missing_from_current_vllm() diff --git a/tests/worker/test_omni_v2_profile_dispatch.py b/tests/worker/test_omni_v2_profile_dispatch.py new file mode 100644 index 000000000..72bcf8d73 --- /dev/null +++ b/tests/worker/test_omni_v2_profile_dispatch.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +from vllm.compilation.cuda_graph import CUDAGraphMode + +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + + +def test_eager_dispatch_uses_eager_without_cudagraph_manager(): + runner = object.__new__(OmniGPUModelRunner) + runner.cudagraph_manager = None + runner.dp_size = 1 + runner.dp_rank = 0 + + batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor( + num_reqs=1, + num_toks=8, + uniform_tok_count=8, + use_eager=True, + ) + + assert batch_desc.cg_mode == CUDAGraphMode.NONE + assert batch_desc.num_tokens == 8 + assert batch_desc.num_reqs == 1 + assert num_tokens_across_dp is None + + +def test_eager_dispatch_syncs_dp_padding(monkeypatch): + runner = object.__new__(OmniGPUModelRunner) + runner.cudagraph_manager = None + runner.dp_size = 2 + runner.dp_rank = 1 + expected_tokens_across_dp = object() + calls = [] + + def sync_dp(cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank): + calls.append((cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank)) + return batch_desc, expected_tokens_across_dp + + import vllm.v1.worker.gpu.dp_utils as dp_utils + + monkeypatch.setattr(dp_utils, "sync_cudagraph_and_dp_padding", sync_dp) + + batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor( + num_reqs=1, + num_toks=8, + uniform_tok_count=8, + use_eager=True, + ) + + assert batch_desc.cg_mode == CUDAGraphMode.NONE + assert batch_desc.num_tokens == 8 + assert batch_desc.num_reqs == 1 + assert num_tokens_across_dp is expected_tokens_across_dp + assert calls == [(None, batch_desc, 8, 1, 8, 2, 1)] + + +def test_non_eager_dispatch_uses_cudagraph_manager(): + runner = object.__new__(OmniGPUModelRunner) + expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1) + calls = [] + runner.cudagraph_manager = SimpleNamespace(dispatch=lambda *args: calls.append(args) or expected) + runner.dp_size = 1 + runner.dp_rank = 0 + + batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor( + num_reqs=1, + num_toks=8, + uniform_tok_count=8, + use_eager=False, + ) + + assert batch_desc is expected + assert calls == [(1, 8, 8)] + assert num_tokens_across_dp is None + + +if __name__ == "__main__": + test_eager_dispatch_uses_eager_without_cudagraph_manager() + test_non_eager_dispatch_uses_cudagraph_manager() + + +def test_skip_compiled_uses_eager_dispatch_without_overwriting_batch_descriptor(): + source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model) + + assert "use_eager=is_profile or skip_compiled" in source + assert "if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs" not in source + + +def test_fullgraph_requires_cudagraph_manager_contract(): + source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model) + + assert "assert self.cudagraph_manager is not None" in source diff --git a/tests/worker/test_omni_v2_shutdown_signature.py b/tests/worker/test_omni_v2_shutdown_signature.py new file mode 100644 index 000000000..f6ca807a7 --- /dev/null +++ b/tests/worker/test_omni_v2_shutdown_signature.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner +from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + + +def test_v2_model_runners_have_shutdown_for_vllm_020_worker_exit(): + for runner_cls in (OmniARModelRunner, OmniGenerationModelRunner, OmniGPUModelRunner): + assert hasattr(runner_cls, "shutdown") + + +if __name__ == "__main__": + test_v2_model_runners_have_shutdown_for_vllm_020_worker_exit() + + +def test_load_model_uses_class_patch_lock(): + import inspect + + source = inspect.getsource(OmniGPUModelRunner.load_model) + assert "_model_state_patch_lock" in source diff --git a/tests/worker_v2/__init__.py b/tests/worker_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/worker_v2/test_init_model_state.py b/tests/worker_v2/test_init_model_state.py new file mode 100644 index 000000000..5c3aabb1c --- /dev/null +++ b/tests/worker_v2/test_init_model_state.py @@ -0,0 +1,83 @@ +"""Unit tests for init_omni_model_state factory dispatch.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.worker_v2.model_states import ( + _OMNI_ARCHITECTURES, + init_omni_model_state, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_vllm_config(architectures): + model_config = SimpleNamespace(architectures=architectures) + return SimpleNamespace(model_config=model_config) + + +@patch("vllm_omni.worker_v2.model_states.omni_model_state.OmniModelState.__init__", return_value=None) +def test_qwen3_omni_dispatches_to_omni_model_state(mock_init): + cfg = _make_vllm_config(["Qwen3OmniMoeForConditionalGeneration"]) + model = MagicMock() + device = torch.device("cpu") + + state = init_omni_model_state(cfg, model, None, device) + + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + + assert isinstance(state, OmniModelState) + + +@patch("vllm_omni.worker_v2.model_states.omni_model_state.OmniModelState.__init__", return_value=None) +def test_mammoth_dispatches_to_omni_model_state(mock_init): + cfg = _make_vllm_config(["MammothModa2ForConditionalGeneration"]) + model = MagicMock() + device = torch.device("cpu") + + state = init_omni_model_state(cfg, model, None, device) + + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + + assert isinstance(state, OmniModelState) + + +@patch("vllm_omni.worker_v2.model_states._upstream_init_model_state") +def test_unknown_arch_delegates_to_upstream(mock_upstream): + mock_upstream.return_value = MagicMock() + cfg = _make_vllm_config(["LlamaForCausalLM"]) + model = MagicMock() + device = torch.device("cpu") + + state = init_omni_model_state(cfg, model, None, device) + + mock_upstream.assert_called_once_with(cfg, model, None, device) + assert state is mock_upstream.return_value + + +def test_omni_architectures_set_contains_expected(): + expected = { + "Qwen3OmniMoeForConditionalGeneration", + "Qwen2_5OmniForConditionalGeneration", + "MammothModa2ForConditionalGeneration", + "MiMoAudioForConditionalGeneration", + "MammothModa2ARForConditionalGeneration", + "Qwen3TTSTalkerForConditionalGeneration", + "Qwen3TTSCode2Wav", + } + assert _OMNI_ARCHITECTURES == expected + + +@patch("vllm_omni.worker_v2.model_states._upstream_init_model_state") +def test_none_architectures_delegates_to_upstream(mock_upstream): + mock_upstream.return_value = MagicMock() + cfg = _make_vllm_config(None) + model = MagicMock() + device = torch.device("cpu") + + init_omni_model_state(cfg, model, None, device) + + mock_upstream.assert_called_once() diff --git a/tests/worker_v2/test_intermediate_buffer.py b/tests/worker_v2/test_intermediate_buffer.py new file mode 100644 index 000000000..f68bd94c5 --- /dev/null +++ b/tests/worker_v2/test_intermediate_buffer.py @@ -0,0 +1,168 @@ +"""Unit tests for OmniIntermediateBuffer lifecycle.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + OmniIntermediateBuffer, + _resolve_additional_information, + _resolve_prompt_embeds, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyInputBatch: + def __init__(self, idx_mapping): + self.idx_mapping_np = idx_mapping + + +def _make_new_req_data(req_id="r1", prompt_embeds=None, mm_features=None, additional_information=None): + ns = SimpleNamespace( + req_id=req_id, + mm_features=mm_features or [], + ) + if prompt_embeds is not None: + ns.prompt_embeds = prompt_embeds + if additional_information is not None: + ns.additional_information = additional_information + return ns + + +# --------------------------------------------------------------- +# _resolve_prompt_embeds +# --------------------------------------------------------------- + + +def test_resolve_prompt_embeds_tensor(): + t = torch.randn(3, 4, device="cpu") + result = _resolve_prompt_embeds(t) + assert isinstance(result, torch.Tensor) + assert result.shape == (3, 4) + assert result.is_contiguous() + + +def test_resolve_prompt_embeds_none(): + assert _resolve_prompt_embeds(None) is None + + +# --------------------------------------------------------------- +# _resolve_additional_information +# --------------------------------------------------------------- + + +def test_resolve_additional_information_dict_passthrough(): + d = {"key": "val"} + assert _resolve_additional_information(d) == d + + +def test_resolve_additional_information_none(): + assert _resolve_additional_information(None) == {} + + +# --------------------------------------------------------------- +# OmniIntermediateBuffer core lifecycle +# --------------------------------------------------------------- + + +def test_add_and_remove(): + buf = OmniIntermediateBuffer(max_num_reqs=4) + req = _make_new_req_data(req_id="r1") + buf.add_request(0, req) + assert buf.buffers[0]["req_id"] == "r1" + + buf.remove_request(0) + assert buf.buffers[0] == {} + + +def test_add_with_prompt_embeds(): + pe = torch.randn(2, 8) + req = _make_new_req_data(req_id="r1", prompt_embeds=pe) + buf = OmniIntermediateBuffer(max_num_reqs=4) + buf.add_request(0, req) + assert "prompt_embeds_cpu" in buf.buffers[0] + assert buf.buffers[0]["prompt_embeds_cpu"].shape == (2, 8) + + +def test_add_with_mm_features(): + features = [torch.randn(3, 4)] + req = _make_new_req_data(req_id="r1", mm_features=features) + buf = OmniIntermediateBuffer(max_num_reqs=4) + buf.add_request(1, req) + assert buf.buffers[1]["mm_features"] is features + + +def test_add_with_additional_information(): + info = {"audio_embeds": torch.randn(5)} + req = _make_new_req_data(req_id="r1", additional_information=info) + buf = OmniIntermediateBuffer(max_num_reqs=4) + buf.add_request(0, req) + assert torch.allclose(buf.buffers[0]["audio_embeds"], info["audio_embeds"]) + + +# --------------------------------------------------------------- +# gather +# --------------------------------------------------------------- + + +def test_gather_returns_batch_order(): + buf = OmniIntermediateBuffer(max_num_reqs=4) + buf.add_request(0, _make_new_req_data(req_id="r0")) + buf.add_request(2, _make_new_req_data(req_id="r2")) + buf.add_request(3, _make_new_req_data(req_id="r3")) + + batch = _DummyInputBatch(idx_mapping=[2, 0, 3]) + gathered = buf.gather(batch) + assert [g["req_id"] for g in gathered] == ["r2", "r0", "r3"] + + +# --------------------------------------------------------------- +# update +# --------------------------------------------------------------- + + +def test_update_merges_tensors_to_cpu(): + buf = OmniIntermediateBuffer(max_num_reqs=2) + buf.add_request(0, _make_new_req_data(req_id="r0")) + + buf.update(0, {"hidden": torch.randn(4)}) + assert "hidden" in buf.buffers[0] + assert buf.buffers[0]["hidden"].device == torch.device("cpu") + + +def test_update_gpu_resident_keys(): + buf = OmniIntermediateBuffer(max_num_reqs=2) + buf.add_request(0, _make_new_req_data(req_id="r0")) + + t = torch.randn(4) + buf.update(0, {"kv": t}, gpu_resident_keys={"kv"}) + assert "kv" in buf.buffers[0] + + +def test_update_list_values(): + buf = OmniIntermediateBuffer(max_num_reqs=2) + buf.add_request(0, _make_new_req_data(req_id="r0")) + + items = [torch.randn(2), "text"] + buf.update(0, {"items": items}) + stored = buf.buffers[0]["items"] + assert isinstance(stored[0], torch.Tensor) + assert stored[1] == "text" + + +def test_update_empty_is_noop(): + buf = OmniIntermediateBuffer(max_num_reqs=2) + buf.add_request(0, _make_new_req_data(req_id="r0")) + before = dict(buf.buffers[0]) + buf.update(0, {}) + assert buf.buffers[0] == before + + +def test_update_accumulates(): + buf = OmniIntermediateBuffer(max_num_reqs=2) + buf.add_request(0, _make_new_req_data(req_id="r0")) + buf.update(0, {"a": torch.tensor([1.0])}) + buf.update(0, {"b": torch.tensor([2.0])}) + assert "a" in buf.buffers[0] and "b" in buf.buffers[0] diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py new file mode 100644 index 000000000..57b3e3c25 --- /dev/null +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -0,0 +1,279 @@ +"""Unit tests for OmniARModelRunner v2.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.data_entry_keys import unflatten_payload +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.omni_ar_model_runner import ( + OmniARModelRunner, + _async_copy_mm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------- +# _build_pooler_output_from_cpu (was _build_pooler_output) +# --------------------------------------------------------------- + + +def test_reconstruct_raw_model_output_preserves_omni_multimodal_outputs(): + hidden = torch.randn(3, 4) + latent = torch.randn(3, 4) + raw = OmniARModelRunner._reconstruct_raw_model_output( + hidden_states=hidden, + multimodal_outputs={"latent": latent}, + aux=None, + ) + + assert isinstance(raw, OmniOutput) + assert raw.text_hidden_states is hidden + assert raw.multimodal_outputs["latent"] is latent + + +def test_reconstruct_raw_model_output_keeps_aux_tuple_without_multimodal_outputs(): + hidden = torch.randn(3, 4) + aux = {"layers": torch.randn(3, 2)} + raw = OmniARModelRunner._reconstruct_raw_model_output( + hidden_states=hidden, + multimodal_outputs=None, + aux=aux, + ) + + assert raw == (hidden, aux) + + +def test_reconstruct_raw_model_output_ignores_empty_multimodal_outputs(): + hidden = torch.randn(3, 4) + raw = OmniARModelRunner._reconstruct_raw_model_output( + hidden_states=hidden, + multimodal_outputs={}, + aux=None, + ) + + assert raw is hidden + + +def test_build_pooler_output_basic(): + """Verify _build_pooler_output_from_cpu slices per-request hidden + mm.""" + hidden = torch.randn(6, 8) + mm = {"audio": torch.randn(6, 2)} + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm, + query_start_loc_np=np.array([0, 3]), + num_scheduled_tokens=np.array([3, 3], dtype=np.int32), + num_reqs=2, + ) + + assert len(pooler) == 2 + assert pooler[0]["hidden"].shape == (3, 8) + assert pooler[1]["hidden"].shape == (3, 8) + assert pooler[0]["audio"].shape == (3, 2) + + +def test_build_pooler_output_empty_mm(): + hidden = torch.randn(4, 8) + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + {}, + query_start_loc_np=np.array([0]), + num_scheduled_tokens=np.array([4], dtype=np.int32), + num_reqs=1, + ) + assert len(pooler) == 1 + assert "hidden" in pooler[0] + assert len(pooler[0]) == 1 + + +# --------------------------------------------------------------- +# _async_copy_mm (was _copy_mm_to_cpu) +# --------------------------------------------------------------- + + +def test_copy_mm_to_cpu_tensor(): + total = 10 + t = torch.randn(10, 4) + result = _async_copy_mm({"feat": t}, total) + assert "feat" in result + assert result["feat"].shape == (10, 4) + assert result["feat"].device == torch.device("cpu") + + +def test_copy_mm_to_cpu_dict(): + total = 10 + d = {"inner": torch.randn(10, 2)} + result = _async_copy_mm({"nested": d}, total) + assert "nested" in result + assert "inner" in result["nested"] + + +def test_copy_mm_to_cpu_list(): + result = _async_copy_mm({"items": [torch.randn(3), "text"]}, 10) + assert "items" in result + assert isinstance(result["items"][0], torch.Tensor) + assert result["items"][1] == "text" + + +def test_copy_mm_to_cpu_empty(): + assert _async_copy_mm({}, 10) == {} + + +# --------------------------------------------------------------- +# Slicing via _build_pooler_output_from_cpu (was _slice_mm_payload) +# --------------------------------------------------------------- + + +def test_slice_mm_payload_tensor(): + hidden = torch.randn(6, 4) + mm_cpu = {"feat": torch.randn(6, 2)} + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm_cpu, + query_start_loc_np=np.array([0, 3]), + num_scheduled_tokens=np.array([3, 3], dtype=np.int32), + num_reqs=2, + ) + assert pooler[0]["feat"].shape == (3, 2) + assert pooler[1]["feat"].shape == (3, 2) + + +def test_slice_mm_payload_list(): + hidden = torch.randn(6, 4) + mm_cpu = {"items": [torch.randn(2), torch.randn(3)]} + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm_cpu, + query_start_loc_np=np.array([0, 3]), + num_scheduled_tokens=np.array([3, 3], dtype=np.int32), + num_reqs=2, + ) + assert isinstance(pooler[0]["items"], torch.Tensor) + assert isinstance(pooler[1]["items"], torch.Tensor) + + +def test_slice_mm_payload_dict(): + hidden = torch.randn(6, 4) + mm_cpu = {"nested": {"a": torch.randn(6, 2)}} + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm_cpu, + query_start_loc_np=np.array([0, 3]), + num_scheduled_tokens=np.array([3, 3], dtype=np.int32), + num_reqs=2, + ) + assert pooler[1]["nested"]["a"].shape == (3, 2) + + +def test_build_pooler_output_flattens_nested_payload_for_msgspec(): + hidden = torch.randn(4, 4) + mm_cpu = {"codes": {"audio": torch.randn(4, 16)}} + + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm_cpu, + query_start_loc_np=np.array([0, 2]), + num_scheduled_tokens=np.array([2, 2], dtype=np.int32), + num_reqs=2, + ) + + assert "codes" not in pooler[0] + assert pooler[0]["codes.audio"].shape == (2, 16) + assert pooler[1]["codes.audio"].shape == (2, 16) + + +def test_build_pooler_output_preserves_qwen3_nested_payload(): + hidden = torch.randn(4, 4) + mm = { + "hidden_states": { + "layers": { + 0: torch.randn(4, 4), + 24: torch.randn(4, 4), + }, + }, + "embed": { + "tts_bos": [torch.randn(1, 1, 4)], + "tts_eos": [torch.randn(1, 1, 4)], + "tts_pad": [torch.randn(1, 1, 4)], + }, + } + + mm_cpu = _async_copy_mm(mm, total_tokens=4) + pooler = OmniARModelRunner._build_pooler_output_from_cpu( + hidden, + mm_cpu, + query_start_loc_np=np.array([0, 2]), + num_scheduled_tokens=np.array([2, 2], dtype=np.int32), + num_reqs=2, + ) + + payload = unflatten_payload(pooler[0]) + assert payload["hidden_states"]["layers"][0].shape == (2, 4) + assert payload["hidden_states"]["layers"][24].shape == (2, 4) + assert payload["embed"]["tts_bos"].shape == (1, 1, 4) + + +def test_clamp_sampling_prompt_token_ids_to_logits_vocab(): + prompt_token_ids = torch.tensor([[3, 99, 152064]]) + input_batch = SimpleNamespace( + vocab_size=152064, + sampling_metadata=SimpleNamespace( + no_penalties=False, + prompt_token_ids=prompt_token_ids, + ), + ) + + OmniARModelRunner._clamp_sampling_prompt_token_ids( + input_batch, + logits_vocab_size=10, + ) + + assert prompt_token_ids.tolist() == [[3, 9, 9]] + + +def test_sample_clamps_prompt_token_ids_from_actual_logits_vocab(): + runner = object.__new__(OmniARModelRunner) + + class Model: + def compute_logits(self, *_args, **_kwargs): + return torch.empty(1, 7) + + runner.model = Model() + prompt_token_ids = torch.tensor([[3, 99, 152064]]) + input_batch = SimpleNamespace( + vocab_size=152064, + sampling_metadata=SimpleNamespace( + no_penalties=False, + prompt_token_ids=prompt_token_ids, + ), + ) + sample_seen_prompt_ids = [] + + def sample(text_hidden, input_batch_arg, grammar_output): + assert grammar_output is None + runner.model.compute_logits(text_hidden) + sample_seen_prompt_ids.append(input_batch_arg.sampling_metadata.prompt_token_ids.clone()) + return "sampler", "num_sampled", "num_rejected" + + runner.sample = sample + + result = runner._sample_with_prompt_token_compat( + torch.empty(1, 4), + input_batch, + None, + ) + + assert result == ("sampler", "num_sampled", "num_rejected") + assert prompt_token_ids.tolist() == [[3, 6, 6]] + assert sample_seen_prompt_ids[0].tolist() == [[3, 6, 6]] + assert runner.model.compute_logits.__func__ is Model.compute_logits + assert "compute_logits" not in runner.model.__dict__ diff --git a/tests/worker_v2/test_omni_forward_compat.py b/tests/worker_v2/test_omni_forward_compat.py new file mode 100644 index 000000000..c5d41337d --- /dev/null +++ b/tests/worker_v2/test_omni_forward_compat.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + + +def _load_add_forward_compat_kwargs(): + root = Path(__file__).resolve().parents[2] + module_path = root / "vllm_omni" / "worker_v2" / "forward_compat.py" + spec = importlib.util.spec_from_file_location("omni_forward_compat", module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.add_forward_compat_kwargs + + +def test_forward_compat_kwargs_include_sampling_metadata_and_sampler(): + add_forward_compat_kwargs = _load_add_forward_compat_kwargs() + model_inputs = {"input_ids": object()} + sampling_metadata = object() + sampler = object() + input_batch = SimpleNamespace( + sampling_metadata=sampling_metadata, + logits_indices=[3, 7], + ) + + add_forward_compat_kwargs(model_inputs, input_batch, sampler) + + assert model_inputs["sampling_metadata"] is sampling_metadata + assert model_inputs["logits_index"] == [3, 7] + assert model_inputs["sampler"] is sampler + + +def test_forward_compat_kwargs_do_not_clobber_explicit_values(): + add_forward_compat_kwargs = _load_add_forward_compat_kwargs() + model_inputs = { + "sampling_metadata": "explicit-metadata", + "logits_index": "explicit-logits", + "sampler": "explicit-sampler", + } + input_batch = SimpleNamespace( + sampling_metadata="batch-metadata", + logits_indices="batch-logits", + ) + + add_forward_compat_kwargs(model_inputs, input_batch, sampler="runner-sampler") + + assert model_inputs["sampling_metadata"] == "explicit-metadata" + assert model_inputs["logits_index"] == "explicit-logits" + assert model_inputs["sampler"] == "explicit-sampler" + + +if __name__ == "__main__": + test_forward_compat_kwargs_include_sampling_metadata_and_sampler() + test_forward_compat_kwargs_do_not_clobber_explicit_values() diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py new file mode 100644 index 000000000..40600d8c0 --- /dev/null +++ b/tests/worker_v2/test_omni_generation_model_runner.py @@ -0,0 +1,260 @@ +"""Tests for OmniGenerationModelRunner.sample_tokens (V2). + +Covers the core pooler_output construction paths via _build_pooler_output: + - OmniOutput with batched tensor multimodal_outputs → per-request slicing + - OmniOutput with list multimodal_outputs → direct mapping (including None) + - OmniOutput with dict scalar values → broadcast to all requests + - None model output → returns None + - Non-dict multimodal_outputs → [None] * num_reqs + - sampled_token_ids always emits empty lists per request (no token sampling) + - req_states.num_computed_tokens updated to prompt_len after sample_tokens +""" + +import unittest +from unittest.mock import MagicMock + +import numpy as np +import torch + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput + +pytestmark = [] + + +class _FakeInputBatch: + """Minimal input batch for sample_tokens.""" + + def __init__(self, num_reqs: int = 1, req_ids: list[str] | None = None): + self.num_reqs = num_reqs + self.req_ids = req_ids or [f"req-{i}" for i in range(num_reqs)] + self.idx_mapping_np = np.arange(num_reqs, dtype=np.int32) + + +class _FakeStagedField: + """Minimal mock for req_states fields that support staged writes.""" + + def __init__(self, data: np.ndarray): + self.np = data + self._staged: list[tuple[int, int]] = [] + + def stage_write_elem(self, idx: int, value: int) -> None: + self._staged.append((idx, value)) + + def apply_write(self) -> None: + for idx, value in self._staged: + self.np[idx] = value + self._staged.clear() + + +class _FakeNpField: + """Minimal mock for req_states fields with .np attribute.""" + + def __init__(self, data: np.ndarray): + self.np = data + + +def _make_omni_output(multimodal_outputs: dict | None = None) -> OmniOutput: + """Create an OmniOutput with given multimodal_outputs.""" + return OmniOutput( + text_hidden_states=torch.zeros(1), + multimodal_outputs=multimodal_outputs, + ) + + +def _make_runner( + model_output, + num_reqs: int = 1, + prompt_len: int = 10, +): + """Build a minimal OmniGenerationModelRunner for sample_tokens testing.""" + from vllm_omni.worker_v2.omni_generation_model_runner import ( + OmniGenerationModelRunner, + ) + + runner = object.__new__(OmniGenerationModelRunner) + runner.device = torch.device("cpu") + + mc = MagicMock() + del mc.eos_token_id + mc.hf_text_config = None + runner.model_config = mc + + runner.postprocess = lambda *a, **kw: None + + input_batch = _FakeInputBatch(num_reqs) + runner._gen_model_output = model_output + runner._gen_input_batch = input_batch + runner._gen_kv_connector_output = None + runner.execute_model_state = None + + req_states = MagicMock() + req_states.prompt_len = _FakeNpField( + np.full(num_reqs, prompt_len, dtype=np.int32), + ) + req_states.num_computed_tokens = _FakeStagedField( + np.zeros(num_reqs, dtype=np.int32), + ) + runner.req_states = req_states + + return runner + + +class TestSampleTokensTensorOutput(unittest.TestCase): + def test_single_request(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"model_outputs": torch.randn(1, 4, 8)}) + runner = _make_runner(output, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert isinstance(result, OmniModelRunnerOutput) + assert len(result.pooler_output) == 1 + assert result.pooler_output[0]["model_outputs"].shape == (4, 8) + + def test_multi_request(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"model_outputs": torch.randn(3, 2, 5)}) + runner = _make_runner(output, num_reqs=3) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 3 + for i in range(3): + assert result.pooler_output[i]["model_outputs"].shape == (2, 5) + + +class TestSampleTokensListOutput(unittest.TestCase): + def test_list_of_tensors(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"model_outputs": [torch.randn(3, 2)]}) + runner = _make_runner(output, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 1 + assert result.pooler_output[0]["model_outputs"].shape == (3, 2) + + def test_list_with_none(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"model_outputs": [None]}) + runner = _make_runner(output, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 1 + assert result.pooler_output[0]["model_outputs"] is None + + +class TestSampleTokensDictOutput(unittest.TestCase): + def test_dict_with_batched_tensor(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"audio": torch.randn(2, 16000), "sr": 24000}) + runner = _make_runner(output, num_reqs=2) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 2 + assert result.pooler_output[0]["audio"].shape == (16000,) + assert result.pooler_output[1]["audio"].shape == (16000,) + assert result.pooler_output[0]["sr"] == 24000 + + def test_dict_with_list_values(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"chunks": [torch.randn(10), torch.randn(20)]}) + runner = _make_runner(output, num_reqs=2) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 2 + assert result.pooler_output[0]["chunks"].shape == (10,) + assert result.pooler_output[1]["chunks"].shape == (20,) + + +class TestSampleTokensNoneOutput(unittest.TestCase): + def test_none_model_output(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + runner = _make_runner(None, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + assert result is None + + +class TestNonDictMultimodalOutputs(unittest.TestCase): + """When multimodal_outputs is None or non-dict, pooler_output is [None]*num_reqs.""" + + def test_none_multimodal_outputs(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output(multimodal_outputs=None) + runner = _make_runner(output, num_reqs=2) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.pooler_output) == 2 + assert result.pooler_output[0] is None + assert result.pooler_output[1] is None + + +class TestSampledTokenIds(unittest.TestCase): + def test_empty_sampled_token_ids_per_request(self): + """Generation models emit empty sampled_token_ids (no token sampling).""" + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output({"model_outputs": torch.randn(3, 2)}) + runner = _make_runner(output, num_reqs=3) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert len(result.sampled_token_ids) == 3 + for ids in result.sampled_token_ids: + assert ids == [] + + +class TestReqStatesUpdate(unittest.TestCase): + """Verify that sample_tokens marks all tokens as computed.""" + + def test_num_computed_tokens_set_to_prompt_len(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + prompt_len = 15 + output = _make_omni_output({"model_outputs": torch.randn(2, 4)}) + runner = _make_runner(output, num_reqs=2, prompt_len=prompt_len) + + OmniGenerationModelRunner.sample_tokens(runner) + + for i in range(2): + assert runner.req_states.num_computed_tokens.np[i] == prompt_len + + +class TestMultimodalOutputsPassthrough(unittest.TestCase): + """multimodal_outputs is a per-request list (tensor-only) on OmniModelRunnerOutput. + + OmniGenerationScheduler indexes it as mm_outputs[req_index], so it must be a + list (not the raw dict). Each entry mirrors the per-request pooler payload. + """ + + def test_multimodal_outputs_on_result(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + mm = {"audio": [torch.randn(10)]} + output = _make_omni_output(mm) + runner = _make_runner(output, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + + assert isinstance(result.multimodal_outputs, list) + assert len(result.multimodal_outputs) == 1 + assert "audio" in result.multimodal_outputs[0] + assert torch.is_tensor(result.multimodal_outputs[0]["audio"]) + + def test_none_multimodal_outputs_becomes_empty_dict(self): + from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner + + output = _make_omni_output(multimodal_outputs=None) + runner = _make_runner(output, num_reqs=1) + result = OmniGenerationModelRunner.sample_tokens(runner) + + # No multimodal data -> one empty dict per request. + assert result.multimodal_outputs == [{}] + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py new file mode 100644 index 000000000..dac092f70 --- /dev/null +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -0,0 +1,137 @@ +"""Unit tests for OmniGPUModelRunner v2 overrides.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner, _needs_capture_tensor_unwrap + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _DummyInputBatch: + def __init__(self, idx_mapping=None): + self.idx_mapping_np = idx_mapping or [] + + +def _make_runner(): + """Create an OmniGPUModelRunner without calling __init__.""" + runner = object.__new__(OmniGPUModelRunner) + runner.model = MagicMock() + runner.req_states = SimpleNamespace(req_id_to_index={"r1": 0, "r2": 1}) + runner.execute_model_state = None + return runner + + +def test_finish_requests_calls_remove_for_finished(): + runner = _make_runner() + mock_state = MagicMock() + runner.model_state = mock_state + + sched_output = SimpleNamespace( + finished_req_ids={"r1"}, + preempted_req_ids=set(), + ) + + with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None): + runner.finish_requests(sched_output) + + mock_state.remove_request.assert_called_once_with(0) + + +def test_finish_requests_calls_remove_for_preempted(): + runner = _make_runner() + mock_state = MagicMock() + runner.model_state = mock_state + + sched_output = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids={"r2"}, + ) + + with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None): + runner.finish_requests(sched_output) + + mock_state.remove_request.assert_called_once_with(1) + + +def test_finish_requests_ignores_unknown_req_ids(): + runner = _make_runner() + mock_state = MagicMock() + runner.model_state = mock_state + + sched_output = SimpleNamespace( + finished_req_ids={"unknown"}, + preempted_req_ids=set(), + ) + + with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None): + runner.finish_requests(sched_output) + + mock_state.remove_request.assert_not_called() + + +def test_finish_requests_handles_both_finished_and_preempted(): + runner = _make_runner() + mock_state = MagicMock() + runner.model_state = mock_state + + sched_output = SimpleNamespace( + finished_req_ids={"r1"}, + preempted_req_ids={"r2"}, + ) + + with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None): + runner.finish_requests(sched_output) + + assert mock_state.remove_request.call_count == 2 + + +def test_thinker_stage_needs_capture_tensor_unwrap(): + assert _needs_capture_tensor_unwrap(SimpleNamespace(model_stage="thinker")) + assert not _needs_capture_tensor_unwrap(SimpleNamespace(model_stage="talker")) + + +def test_capture_model_unwraps_tuple_outputs(): + runner = object.__new__(OmniGPUModelRunner) + hidden = torch.ones(1, 2) + + def original_forward(): + return hidden, {"layers": {}} + + runner.model = SimpleNamespace(forward=original_forward) + runner._model_returns_tuple = True + runner._exclude_full_graph = False + + def capture_model(_self): + assert torch.equal(runner.model.forward(), hidden) + return 3 + + with patch.object(type(runner).__bases__[0], "capture_model", capture_model): + assert runner.capture_model() == 3 + + assert runner.model.forward is original_forward + + +def test_capture_model_unwraps_omni_outputs(): + runner = object.__new__(OmniGPUModelRunner) + hidden = torch.ones(1, 2) + + def original_forward(): + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + runner.model = SimpleNamespace(forward=original_forward) + runner._model_returns_tuple = True + runner._exclude_full_graph = False + + def capture_model(_self): + assert torch.equal(runner.model.forward(), hidden) + return 5 + + with patch.object(type(runner).__bases__[0], "capture_model", capture_model): + assert runner.capture_model() == 5 + + assert runner.model.forward is original_forward diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py new file mode 100644 index 000000000..34f33250d --- /dev/null +++ b/tests/worker_v2/test_omni_model_state.py @@ -0,0 +1,334 @@ +"""Unit tests for OmniModelState core methods and plugin dispatch.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.engine.serialization import serialize_additional_information +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.model_states.plugin import OmniModelStatePlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------- +# Stubs +# --------------------------------------------------------------- + + +class _DummyInputBatch: + def __init__(self, indices): + self.idx_mapping_np = indices + self.num_reqs = len(indices) + self.num_scheduled_tokens = [1] * len(indices) + self.query_start_loc_np = list(range(len(indices))) + + +class _DummyReqState: + pass + + +class _SpyPlugin(OmniModelStatePlugin): + """Plugin that records calls for verification.""" + + def __init__(self): + self.add_calls: list = [] + self.remove_calls: list = [] + self.prepare_calls: list = [] + self.postprocess_calls: list = [] + + def on_add_request(self, req_index, new_req_data): + self.add_calls.append((req_index, new_req_data)) + + def on_remove_request(self, req_index): + self.remove_calls.append(req_index) + + def prepare_extra_inputs(self, input_batch, req_states): + self.prepare_calls.append(True) + return {"plugin_key": "plugin_value"} + + def postprocess(self, text_hidden, multimodal_outputs, input_batch, req_states): + self.postprocess_calls.append(True) + return text_hidden, multimodal_outputs + + +def _make_state( + max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False, plugins=None +): + """Create an OmniModelState without calling real __init__.""" + state = object.__new__(OmniModelState) + + # Minimal model mock + model = MagicMock() + model.has_preprocess = has_preprocess + model.has_postprocess = has_postprocess + model.have_multimodal_outputs = have_multimodal_outputs + model.preprocess_in_forward = False + model.get_omni_plugins = MagicMock(return_value=[]) + state.model = model + + # scheduler_config mock + state.scheduler_config = SimpleNamespace(max_num_seqs=max_num_reqs) + + # Skip DefaultModelState.__init__ side effects + state.has_preprocess = has_preprocess + state.has_postprocess = has_postprocess + state.have_multimodal_outputs = have_multimodal_outputs + state.plugins = plugins or [] + + from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + OmniIntermediateBuffer, + ) + + state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) + state._static_inputs_embeds = None + return state + + +def _make_new_req_data(req_id="r1"): + return SimpleNamespace( + req_id=req_id, + mm_features=[], + ) + + +# --------------------------------------------------------------- +# add_request / remove_request +# --------------------------------------------------------------- + + +def test_add_request_populates_buffer(): + state = _make_state() + req = _make_new_req_data("r1") + + with patch.object(type(state).__bases__[0], "add_request", return_value=None): + state.add_request(0, req) + + assert state.intermediate_buffer.buffers[0]["req_id"] == "r1" + + +def test_add_request_dispatches_to_plugins(): + plugin = _SpyPlugin() + state = _make_state(plugins=[plugin]) + req = _make_new_req_data("r1") + + with patch.object(type(state).__bases__[0], "add_request", return_value=None): + state.add_request(0, req) + + assert len(plugin.add_calls) == 1 + assert plugin.add_calls[0][0] == 0 + + +def test_add_request_unflattens_serialized_additional_information(): + state = _make_state() + hidden = torch.randn(2, 4) + embeds = torch.randn(3, 4) + req = SimpleNamespace( + req_id="r1", + mm_features=[], + additional_information=serialize_additional_information( + { + "hidden_states": {"output": hidden}, + "embed": {"prefill": embeds}, + "ids": {"prompt": [1, 2, 3], "output": [4]}, + } + ), + ) + + with patch.object(type(state).__bases__[0], "add_request", return_value=None): + state.add_request(0, req) + + buf = state.intermediate_buffer.buffers[0] + assert torch.equal(buf["hidden_states"]["output"], hidden) + assert torch.equal(buf["embed"]["prefill"], embeds) + assert buf["ids"]["prompt"] == [1, 2, 3] + assert buf["ids"]["output"] == [4] + + +def test_remove_request_clears_buffer(): + state = _make_state() + req = _make_new_req_data("r1") + + with patch.object(type(state).__bases__[0], "add_request", return_value=None): + state.add_request(0, req) + + state.remove_request(0) + assert state.intermediate_buffer.buffers[0] == {} + + +def test_remove_request_dispatches_to_plugins(): + plugin = _SpyPlugin() + state = _make_state(plugins=[plugin]) + state.remove_request(2) + assert plugin.remove_calls == [2] + + +def test_intermediate_buffer_update_merges_nested_and_tuple_keys(): + state = _make_state() + trailing = torch.randn(2, 4) + last = torch.randn(4) + codes = torch.ones(1, 16, dtype=torch.long) + state.intermediate_buffer.buffers[0] = { + "req_id": "r1", + "hidden_states": {"trailing_text": trailing}, + } + + state.intermediate_buffer.update( + 0, + { + ("codes", "audio"): codes, + "hidden_states": {"last": last}, + }, + {("codes", "audio"), ("hidden_states", "last")}, + ) + + buf = state.intermediate_buffer.buffers[0] + assert ("codes", "audio") not in buf + assert torch.equal(buf["codes"]["audio"], codes) + assert torch.equal(buf["hidden_states"]["trailing_text"], trailing) + assert torch.equal(buf["hidden_states"]["last"], last) + + +# --------------------------------------------------------------- +# prepare_inputs +# --------------------------------------------------------------- + + +def test_prepare_inputs_injects_buffer_and_runtime_info(): + state = _make_state() + req = _make_new_req_data("r1") + + with patch.object(type(state).__bases__[0], "add_request", return_value=None): + state.add_request(0, req) + + batch = _DummyInputBatch([0]) + + with patch.object(type(state).__bases__[0], "prepare_inputs", return_value={}): + result = state.prepare_inputs(batch, _DummyReqState()) + + assert "model_intermediate_buffer" in result + assert "runtime_additional_information" in result + assert len(result["model_intermediate_buffer"]) == 1 + assert result["model_intermediate_buffer"][0]["req_id"] == "r1" + + +def test_prepare_inputs_merges_plugin_extra(): + plugin = _SpyPlugin() + state = _make_state(plugins=[plugin]) + + batch = _DummyInputBatch([]) + + with patch.object(type(state).__bases__[0], "prepare_inputs", return_value={"base": True}): + result = state.prepare_inputs(batch, _DummyReqState()) + + assert result["base"] is True + assert result["plugin_key"] == "plugin_value" + + +# --------------------------------------------------------------- +# prepare_dummy_inputs +# --------------------------------------------------------------- + + +def test_prepare_dummy_inputs_has_buffer_keys(): + state = _make_state() + with patch.object(type(state).__bases__[0], "prepare_dummy_inputs", return_value={}): + result = state.prepare_dummy_inputs(num_reqs=2, num_tokens=16) + + assert "model_intermediate_buffer" in result + assert len(result["model_intermediate_buffer"]) == 2 + + +# --------------------------------------------------------------- +# postprocess_model_output +# --------------------------------------------------------------- + + +def test_postprocess_omni_output_with_multimodal(): + state = _make_state(have_multimodal_outputs=True) + hidden = torch.randn(4, 8) + mm = {"audio": torch.randn(4, 2)} + omni = OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + batch = _DummyInputBatch([0]) + th, mo = state.postprocess_model_output(omni, batch, _DummyReqState()) + assert th is hidden + assert mo is mm + + +def test_postprocess_tuple_output(): + state = _make_state(have_multimodal_outputs=False) + del state.model.make_omni_output + hidden = torch.randn(4, 8) + + batch = _DummyInputBatch([0]) + th, mo = state.postprocess_model_output((hidden, {}), batch, _DummyReqState()) + assert th is hidden + assert mo == {} + + +def test_postprocess_raw_tensor(): + state = _make_state() + del state.model.make_omni_output + hidden = torch.randn(4, 8) + + batch = _DummyInputBatch([0]) + th, mo = state.postprocess_model_output(hidden, batch, _DummyReqState()) + assert th is hidden + assert mo == {} + + +def test_postprocess_calls_make_omni_output_when_not_omni(): + """When model output is not OmniOutput but model has make_omni_output, it should be called.""" + state = _make_state(have_multimodal_outputs=True) + hidden = torch.randn(4, 8) + mm = {"audio": torch.randn(4, 2)} + expected_omni = OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + state.model.make_omni_output = MagicMock(return_value=expected_omni) + + batch = _DummyInputBatch([0]) + th, mo = state.postprocess_model_output((hidden, {}), batch, _DummyReqState()) + state.model.make_omni_output.assert_called_once() + assert th is hidden + + +def test_postprocess_dispatches_to_plugins(): + plugin = _SpyPlugin() + state = _make_state(plugins=[plugin]) + hidden = torch.randn(4, 8) + + batch = _DummyInputBatch([0]) + state.postprocess_model_output(hidden, batch, _DummyReqState()) + assert len(plugin.postprocess_calls) == 1 + + +def test_run_postprocess_does_not_duplicate_hidden_states_kwarg(): + state = _make_state(has_postprocess=True) + state.model.gpu_resident_buffer_keys = set() + seen = {} + trailing = torch.ones(1, 2) + + def postprocess(hidden_states, **info): + seen["hidden_states"] = hidden_states + seen["info"] = info + return {"hidden_states": {"last": hidden_states[-1].detach()}} + + state.model.postprocess = postprocess + state.intermediate_buffer.buffers[0] = { + "req_id": "r1", + "hidden_states": {"trailing_text": trailing}, + "meta": {"codec_streaming": True}, + } + + hidden = torch.randn(1, 2) + state.run_postprocess(hidden, _DummyInputBatch([0])) + + assert torch.equal(seen["hidden_states"], hidden) + assert "hidden_states" not in seen["info"] + assert seen["info"]["meta"] == {"codec_streaming": True} + assert torch.equal(state.intermediate_buffer.buffers[0]["hidden_states"]["trailing_text"], trailing) + assert torch.equal(state.intermediate_buffer.buffers[0]["hidden_states"]["last"], hidden[-1]) diff --git a/tests/worker_v2/test_omni_model_state_plugin.py b/tests/worker_v2/test_omni_model_state_plugin.py new file mode 100644 index 000000000..272004e4b --- /dev/null +++ b/tests/worker_v2/test_omni_model_state_plugin.py @@ -0,0 +1,71 @@ +"""Unit tests for OmniModelStatePlugin interface correctness.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from vllm_omni.worker_v2.model_states.plugin import OmniModelStatePlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _ConcretePlugin(OmniModelStatePlugin): + """Minimal concrete plugin for testing the ABC.""" + + def __init__(self): + self.added: list[tuple[int, Any]] = [] + self.removed: list[int] = [] + + def on_add_request(self, req_index, new_req_data): + self.added.append((req_index, new_req_data)) + + def on_remove_request(self, req_index): + self.removed.append(req_index) + + def postprocess(self, text_hidden, multimodal_outputs, input_batch, req_states): + return text_hidden, multimodal_outputs + + +def test_concrete_plugin_instantiates(): + plugin = _ConcretePlugin() + assert isinstance(plugin, OmniModelStatePlugin) + + +def test_on_add_request_called(): + plugin = _ConcretePlugin() + plugin.on_add_request(0, SimpleNamespace(req_id="r1")) + assert len(plugin.added) == 1 + assert plugin.added[0][0] == 0 + + +def test_on_remove_request_called(): + plugin = _ConcretePlugin() + plugin.on_remove_request(3) + assert plugin.removed == [3] + + +def test_prepare_extra_inputs_default_empty(): + plugin = _ConcretePlugin() + result = plugin.prepare_extra_inputs(SimpleNamespace(), SimpleNamespace()) + assert result == {} + + +def test_dummy_run_default_noop(): + plugin = _ConcretePlugin() + plugin.dummy_run(None, 128) + + +def test_postprocess_returns_inputs(): + plugin = _ConcretePlugin() + hidden = torch.randn(4, 8) + mm = {"audio": torch.randn(4, 2)} + out_h, out_mm = plugin.postprocess(hidden, mm, None, None) + assert out_h is hidden + assert out_mm is mm + + +def test_abc_cannot_instantiate_directly(): + with pytest.raises(TypeError): + OmniModelStatePlugin() diff --git a/vllm_omni/compat.py b/vllm_omni/compat.py new file mode 100644 index 000000000..8331e4f56 --- /dev/null +++ b/vllm_omni/compat.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import inspect +from typing import Any + + +def make_filtered_namedtuple(cls, known_extra_fields: set[str] | None = None, **kwargs: Any): + fields = set(cls._fields) + unknown = set(kwargs) - fields - (known_extra_fields or set()) + values = {key: value for key, value in kwargs.items() if key in fields} + return cls(**values), unknown + + +def make_filtered_call(cls, known_extra_fields: set[str] | None = None, **kwargs: Any): + params = set(inspect.signature(cls).parameters) + unknown = set(kwargs) - params - (known_extra_fields or set()) + values = {key: value for key, value in kwargs.items() if key in params} + return cls(**values), unknown diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 984844bdb..4c12d2d5a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -19,6 +19,7 @@ from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.compat import make_filtered_call from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, @@ -28,12 +29,40 @@ from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) -from vllm_omni.engine import OmniEngineCoreOutput -from vllm_omni.engine.serialization import deserialize_additional_information +from vllm_omni.engine.serialization import ( + deserialize_additional_information, + serialize_additional_information, +) from vllm_omni.outputs import OmniConnectorOutput logger = init_logger(__name__) +_KNOWN_ENGINE_CORE_OUTPUT_COMPAT_FIELDS = { + "num_cached_tokens", + "num_external_computed_tokens", + "is_segment_finished", + "new_prompt_len_snapshot", +} + + +def _make_engine_core_output(**kwargs): + output, unknown = make_filtered_call( + EngineCoreOutput, + known_extra_fields=_KNOWN_ENGINE_CORE_OUTPUT_COMPAT_FIELDS, + **kwargs, + ) + if unknown: + logger.warning("Unknown fields passed to EngineCoreOutput: %s", sorted(unknown)) + return output + + +def _get_request_num_cached_tokens(request) -> int: + return max(getattr(request, "num_cached_tokens", 0), 0) + + +def _get_request_num_external_computed_tokens(request) -> int: + return max(getattr(request, "num_external_computed_tokens", 0), 0) + @dataclass class KVCacheTransferData: @@ -249,6 +278,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): req_id=nr.req_id, external_req_id=(getattr(request, "external_req_id", None) if request else None), prompt_token_ids=nr.prompt_token_ids, + prefill_token_ids=getattr(nr, "prefill_token_ids", None), mm_features=nr.mm_features, sampling_params=nr.sampling_params, pooling_params=nr.pooling_params, @@ -328,7 +358,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): req = self.requests.get(req_id) if req is not None and not req.is_finished(): outputs[req.client_index].append( - OmniEngineCoreOutput( + _make_engine_core_output( request_id=req_id, new_token_ids=[], kv_transfer_params={"kv_ready": True}, @@ -466,21 +496,33 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped: + # vLLM 0.23 strictly decodes EngineCoreOutput.pooling_output as a + # torch.Tensor. MR V2 runners hand off a per-request dict here, so + # serialize it onto the bytes channel; StagePool._poll_stage_raw + # rehydrates it. Tensor pooling outputs ride pooling_output as before. + pooling_output_wire = pooler_output + pooling_output_payload = None + if isinstance(pooler_output, dict): + pooling_output_payload = serialize_additional_information(pooler_output) + pooling_output_wire = None # Add EngineCoreOutput for this Request. outputs[request.client_index].append( - OmniEngineCoreOutput( + _make_engine_core_output( request_id=req_id, new_token_ids=new_token_ids, finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_wire, + pooling_output_payload=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, events=request.take_events(), prefill_stats=request.take_prefill_stats(), kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, + num_cached_tokens=_get_request_num_cached_tokens(request), + num_external_computed_tokens=_get_request_num_external_computed_tokens(request), routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, is_segment_finished=is_segment_finished, @@ -507,12 +549,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) for request in requests: outputs[request.client_index].append( - OmniEngineCoreOutput( + _make_engine_core_output( request_id=request.request_id, new_token_ids=[], finish_reason=request.get_finished_reason(), events=request.take_events(), trace_headers=request.trace_headers, + num_cached_tokens=_get_request_num_cached_tokens(request), ) ) if self.chunk_transfer_adapter is not None: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 2e99c8e1a..81e328be7 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import time from collections import defaultdict from typing import Any @@ -23,6 +24,7 @@ from vllm.v1.metrics.perf import PerfStats from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.compat import make_filtered_call from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, @@ -34,10 +36,38 @@ from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapt OmniChunkTransferAdapter, ) from vllm_omni.engine import OmniEngineCoreOutput +from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput logger = init_logger(__name__) +VLLM_OMNI_USE_V2_RUNNER = bool(int(os.environ.get("VLLM_OMNI_USE_V2_RUNNER", "0"))) + +_KNOWN_ENGINE_CORE_OUTPUT_COMPAT_FIELDS = { + "num_cached_tokens", + "num_external_computed_tokens", +} + + +def _make_engine_core_output(**kwargs): + output, unknown = make_filtered_call( + EngineCoreOutput, + known_extra_fields=_KNOWN_ENGINE_CORE_OUTPUT_COMPAT_FIELDS, + **kwargs, + ) + if unknown: + logger.warning("Unknown fields passed to EngineCoreOutput: %s", sorted(unknown)) + return output + + +def _get_request_num_cached_tokens(request) -> int: + return max(getattr(request, "num_cached_tokens", 0), 0) + + +def _set_request_num_cached_tokens_if_present(request, value: int) -> None: + if hasattr(request, "num_cached_tokens"): + request.num_cached_tokens = value + class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): def __init__(self, *args, **kwargs): @@ -314,6 +344,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): req_id=nr.req_id, external_req_id=(getattr(request, "external_req_id", None) if request else None), prompt_token_ids=nr.prompt_token_ids, + prefill_token_ids=nr.prompt_token_ids, mm_features=nr.mm_features, sampling_params=nr.sampling_params, pooling_params=nr.pooling_params, @@ -561,21 +592,36 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped: + num_cached = _get_request_num_cached_tokens(request) + if num_cached < 0: + logger.warning("Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, req_id) + num_cached = 0 + # vLLM 0.23 strictly decodes EngineCoreOutput.pooling_output as a + # torch.Tensor; a dict-shaped MR V2 payload rides the bytes channel + # instead (rehydrated in StagePool._poll_stage_raw). See OmniARScheduler. + pooling_output_wire = pooler_output + pooling_output_payload = None + if isinstance(pooler_output, dict): + pooling_output_payload = serialize_additional_information(pooler_output) + pooling_output_wire = None # Add EngineCoreOutput for this Request. outputs[request.client_index].append( - OmniEngineCoreOutput( + _make_engine_core_output( request_id=req_id, new_token_ids=new_token_ids, finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_wire, + pooling_output_payload=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, events=request.take_events(), prefill_stats=request.take_prefill_stats(), kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, + num_cached_tokens=num_cached, + num_external_computed_tokens=getattr(request, "num_external_computed_tokens", 0), routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) @@ -627,13 +673,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids] self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) for request in requests: + num_cached = _get_request_num_cached_tokens(request) + if num_cached < 0: + logger.warning( + "Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, request.request_id + ) + num_cached = 0 outputs[request.client_index].append( - OmniEngineCoreOutput( + _make_engine_core_output( request_id=request.request_id, new_token_ids=[], finish_reason=request.get_finished_reason(), events=request.take_events(), trace_headers=request.trace_headers, + num_cached_tokens=num_cached, ) ) if self.chunk_transfer_adapter is not None: diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index eec3e5c33..20f670038 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -338,18 +338,31 @@ def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry: from vllm_omni.engine import AdditionalInformationEntry t_cpu = t.detach().to("cpu").contiguous() + try: + tensor_data = t_cpu.numpy().tobytes() + except TypeError: + # numpy has no equivalent for some torch dtypes (bfloat16, float8_*). + # Fall back to a raw byte view; _deserialize_tensor reconstructs via torch. + tensor_data = t_cpu.view(torch.uint8).numpy().tobytes() return AdditionalInformationEntry( - tensor_data=t_cpu.numpy().tobytes(), + tensor_data=tensor_data, tensor_shape=list(t_cpu.shape), tensor_dtype=_dtype_to_name(t_cpu.dtype), ) def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor: - dt = np.dtype(entry.tensor_dtype or "float32") - arr = np.frombuffer(entry.tensor_data, dtype=dt) # type: ignore[arg-type] - arr = arr.reshape(entry.tensor_shape) - return torch.from_numpy(arr.copy()) + name = entry.tensor_dtype or "float32" + try: + dt = np.dtype(name) + arr = np.frombuffer(entry.tensor_data, dtype=dt) # type: ignore[arg-type] + return torch.from_numpy(arr.copy().reshape(entry.tensor_shape)) + except TypeError: + # numpy can't represent this dtype (bfloat16, float8_*): reconstruct from + # the raw byte view written by _serialize_tensor. + torch_dtype = getattr(torch, name) + flat = torch.frombuffer(bytearray(entry.tensor_data), dtype=torch_dtype) + return flat.reshape(entry.tensor_shape).clone() def serialize_payload( diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 89bb41908..b53d686cc 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -121,6 +121,13 @@ class OmniEngineCoreOutput(EngineCoreOutput): # pooling_output is inherited from EngineCoreOutput as torch.Tensor | None # and retains its original vLLM semantics for pooling/embedding tasks. multimodal_output: dict[str, torch.Tensor] | None = None + # Bytes-serialized carrier for a dict-shaped pooling_output (MR V2 cross-stage + # hidden/codec handoff). vLLM 0.23 strictly decodes ``pooling_output`` as a + # ``torch.Tensor``, so a per-request dict payload cannot ride that field. The + # producer (OmniARScheduler) sets ``pooling_output=None`` and serializes the + # dict here; the consumer (StagePool._poll_stage_raw) rehydrates it back into + # ``pooling_output`` after decode. Tensor pooling outputs are unaffected. + pooling_output_payload: AdditionalInformationPayload | None = None # Finished flag for streaming input segment is_segment_finished: bool | None = False # Streaming update prompt length diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index ac6140291..d0828a145 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1532,6 +1532,28 @@ class AsyncOmniEngine: except Exception: pass + # ── Release CuMem allocator memory pool ────────────────────────────── + # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory + # Management) allocator holds model weights in a singleton memory pool + # that lives in the parent process. Killing the engine-core subprocess + # does NOT release this pool — the weights stay resident on the GPU + # and can cause CUDA OOM for subsequent engine instances (especially + # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). + # + # Discard mode (level=2) is correct at shutdown: there is no benefit to + # keeping a CPU backup when the engine is being torn down. + try: + from vllm.device_allocator.cumem import CuMemAllocator, cumem_available + + if cumem_available: + allocator = CuMemAllocator.get_instance() + # Sleep at level 2 discards all pool memory from the GPU + # without creating CPU backups — cheapest and fastest. + allocator.sleep() + logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") + except Exception: + pass + def _try_shutdown(self, *args, **kwargs) -> None: try: self.shutdown() diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 4e1f3d5a5..ede23ee55 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -23,6 +23,8 @@ from vllm_omni.engine.stage_client import ( StagePoolDiffusionClient, StagePoolLLMClient, ) +from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.metrics import definitions as defs from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics @@ -1020,8 +1022,30 @@ class StagePool: outputs = await client.get_output_async() if not outputs.outputs: return None + self._rehydrate_pooling_output_payloads(outputs) return outputs + @staticmethod + def _rehydrate_pooling_output_payloads(outputs: EngineCoreOutputs) -> None: + """Restore dict-shaped pooling_output from its bytes carrier (MR V2). + + vLLM 0.23 decodes EngineCoreOutput.pooling_output as a torch.Tensor, so MR + V2 runners ship the per-request dict handoff via ``pooling_output_payload`` + with ``pooling_output=None``. Decode it back here so downstream stage-input + processors see ``pooling_output`` exactly as on vLLM <= 0.22. + """ + for eco in outputs.outputs: + payload = getattr(eco, "pooling_output_payload", None) + if payload is None: + continue + if getattr(eco, "pooling_output", None) is None: + # deserialize_* returns the nested OmniPayload form; the producer + # serialized a flat (dotted-key) pooler dict, and downstream + # consumers (e.g. talker2code2wav_full_payload) read flat keys like + # "codes.audio". Re-flatten to restore the exact on-wire shape. + eco.pooling_output = flatten_payload(deserialize_additional_information(payload)) + eco.pooling_output_payload = None + async def process_llm_raw_outputs( self, replica_id: int, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index fc22493fc..b9e398e12 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -7,6 +7,7 @@ from collections.abc import Sequence from typing import Any, Literal import huggingface_hub +from vllm import envs from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError @@ -29,6 +30,10 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingNamespace +_omni_use_v2 = os.environ.get("VLLM_OMNI_USE_V2_RUNNER", "0") +os.environ.setdefault("VLLM_USE_V2_MODEL_RUNNER", _omni_use_v2) +envs.VLLM_USE_V2_MODEL_RUNNER = _omni_use_v2 == "1" + logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index afd0a92a5..96e9726fd 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -43,7 +43,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": True, - "stop_token_ids": [8294], + "stop_token_ids": [8292, 8294], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 4a144bff0..a3622b597 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -320,6 +320,12 @@ class Qwen3OmniMoeForConditionalGeneration( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec] | None = None, + # V1 runner passes the full multimodal RoPE inputs explicitly, while + # vLLM 0.20 calls this method with only input_tokens/mm_features. + hf_config: object | None = None, + image_grid_thw: list[list[int]] | torch.Tensor | None = None, + video_grid_thw: list[list[int]] | torch.Tensor | None = None, + second_per_grid_ts: list[float] | None = None, **kwargs: object, ) -> tuple[torch.Tensor, int]: if self.model_stage == "thinker": 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 44aaaf4d3..31feb7713 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 @@ -30,7 +30,6 @@ from typing import Any, cast import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F from packaging.version import Version from transformers import PretrainedConfig from transformers import __version__ as TRANSFORMERS_VERSION @@ -163,66 +162,42 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw) if isinstance(grid_thw, torch.Tensor): - grid_thw_tensor = grid_thw.to(self.device) + grid_thw_np = grid_thw.cpu().numpy().astype(np.int32) else: - grid_thw_tensor = torch.as_tensor(grid_thw, dtype=torch.int32, device=self.device) + grid_thw_np = np.array(grid_thw, dtype=np.int32) - try: - cu_seqlens = torch.repeat_interleave( - grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2], - grid_thw_tensor[:, 0], - ).cumsum( - dim=0, - dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - except RuntimeError: - logger.warning( - "torch.repeat_interleave not executable, switching to vectorized searchsorted implementation." - ) - repeat_counts = grid_thw_tensor[:, 0] - values = grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2] - repeat_cumsum = repeat_counts.cumsum(0) - total_items = repeat_cumsum[-1].item() - indices = torch.searchsorted( - repeat_cumsum, - torch.arange(total_items, device=grid_thw_tensor.device), - right=True, - ) - cu_seqlens = values[indices].cumsum( - dim=0, - dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32) + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens]) - hidden_states = hidden_states.unsqueeze(1) - rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device) - rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device) - max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) - - grid_thw_np = grid_thw_tensor.cpu().numpy().astype(np.int32) - cu_seqlens_np = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum( - axis=0, dtype=np.int32 + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(self.attn_backend, cu_seqlens, self.device) + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens), + dtype=torch.int32, + device=self.device, ) - cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np]) - sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( self.attn_backend, - cu_seqlens_np, + cu_seqlens, + self.hidden_size, + self.tp_size, self.device, ) + hidden_states = hidden_states.unsqueeze(1) + hidden_states_list = [] deepstack_visual_indexes = self.deepstack_visual_indexes for layer_num, blk in enumerate(self.blocks): - hidden_states = blk( - hidden_states, + hidden_states = hidden_states + blk.attn( + blk.norm1(hidden_states), cu_seqlens=cu_seqlens, rotary_pos_emb_cos=rotary_pos_emb_cos, rotary_pos_emb_sin=rotary_pos_emb_sin, max_seqlen=max_seqlen, sequence_lengths=sequence_lengths, ) + hidden_states = hidden_states + blk.mlp(blk.norm2(hidden_states)) if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes: hidden_states_list.append(hidden_states) @@ -784,13 +759,19 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( mm_item_counts, ) else: - if use_audio_in_video and "audio" in mm_prompt_updates: + if use_audio_in_video: + # When use_audio_in_video=True, audio is extracted from video and embedded + # in the video placeholder tokens. We should: + # 1. Filter out audio from prompt updates (audio has no separate placeholder) + # 2. Apply remaining updates (video, image, etc.) + # 3. Derive audio placeholders from video placeholders filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, filtered_updates, ) - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + # Derive audio placeholders from video placeholders + mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_item_counts) else: prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, @@ -909,7 +890,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( def get_replacement_qwen2_use_audio_in_video(item_idx: int): nonlocal audio_in_video_item_idx - audio_num_features = audio_output_lengths[audio_in_video_item_idx] + audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx] video_grid_thw = out_mm_data["video_grid_thw"][item_idx] audio_in_video_item_idx += 1 @@ -955,17 +936,27 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( def _derive_audio_from_video_placeholders( self, placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], - mm_prompt_updates: MultiModalPromptUpdates, + mm_item_counts: Mapping[str, int], ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: """ Helper to derive audio placeholders from video placeholders when use_audio_in_video=True. + + In use_audio_in_video mode, audio is extracted from video and embedded + within the video placeholder tokens. This function creates audio placeholder + info by extracting the audio token positions from video placeholders. + + Args: + placeholders: Current placeholders (should contain "video") + mm_item_counts: Counts of multimodal items from mm_items.get_all_counts() """ if "video" not in placeholders: return placeholders + # Validate audio and video counts match + # In use_audio_in_video mode, audio count comes from mm_items (extracted from video) num_videos = len(placeholders["video"]) - num_audios = len(mm_prompt_updates.get("audio", [])) + num_audios = mm_item_counts.get("audio", 0) if num_audios != num_videos: raise ValueError( f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}" @@ -1285,7 +1276,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( return [] # The result multimodal_embeddings is tuple of tensors, with each - # tensor corresponding to a multimodal data item (image or video). + # tensor correspoending to a multimodal data item (image or video). multimodal_embeddings: tuple[torch.Tensor, ...] = () # NOTE: It is important to iterate over the keys in this dictionary @@ -1309,6 +1300,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: inputs_embeds = self._embed_text_input_ids( input_ids, diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index 28ca4f362..ad8061269 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -536,6 +536,10 @@ class Qwen3TTSConfig(PretrainedConfig): if any("Code2Wav" in str(a) for a in archs): if hasattr(config, "rope_parameters"): delattr(config, "rope_parameters") + # Code2Wav doesn't use position embeddings, but vLLM validates + # max_model_len <= max_position_embeddings. Set a large value + # to allow the pipeline.yaml max_model_len (65536) to pass validation. + config.max_position_embeddings = 131072 return config diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index d624b4a49..ddb6b6252 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -800,6 +800,14 @@ def thinker2talker_token_only( # ========================= +def _eof_payload(transfer_manager: Any) -> OmniPayloadStruct: + """Return an EOF marker payload when the request is finished with no codes.""" + return { + "codes": {"audio": []}, + "meta": {"left_context_size": 0, "finished": torch.tensor(True)}, + } + + def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], @@ -810,18 +818,28 @@ def talker2code2wav_async_chunk( Multimodal output version. """ if not isinstance(multimodal_output, Mapping): + if is_finished: + return _eof_payload(transfer_manager) return None talker_codes = multimodal_output.get("codes", {}) if not isinstance(talker_codes, dict): + if is_finished: + return _eof_payload(transfer_manager) return None code_predictor_codes = talker_codes.get("audio") if code_predictor_codes is None: + if is_finished: + return _eof_payload(transfer_manager) return None if code_predictor_codes.numel() == 0: + if is_finished: + return _eof_payload(transfer_manager) return None if not code_predictor_codes.any(): + if is_finished: + return _eof_payload(transfer_manager) return None connector = getattr(transfer_manager, "connector", None) @@ -839,6 +857,8 @@ def talker2code2wav_async_chunk( first_codebook = int(code_predictor_codes[0, 0].item()) if first_codebook in stop_token_ids: logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) + if is_finished: + return _eof_payload(transfer_manager) return None request_id = request.external_req_id diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 08551ad69..47b862214 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -19,6 +19,8 @@ from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) +VLLM_OMNI_USE_V2_RUNNER = bool(int(os.environ.get("VLLM_OMNI_USE_V2_RUNNER", "0"))) + class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU worker for autoregressive omni model stages. @@ -59,7 +61,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): f"({visible_device_count})." ) self.device = torch.device(f"cuda:{self.local_rank}") - torch.accelerator.set_device_index(self.device) + current_platform.set_device(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) @@ -94,18 +96,63 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) - if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False + if VLLM_OMNI_USE_V2_RUNNER or self.use_v2_model_runner: + from vllm_omni.worker_v2.omni_ar_model_runner import ( + OmniARModelRunner, + ) - # Construct the model runner - self.model_runner = GPUARModelRunner(self.vllm_config, self.device) + logger.info("Using MR v2 OmniARModelRunner for omni AR stage.") + self.use_v2_model_runner = True + self.model_runner = OmniARModelRunner(self.vllm_config, self.device) + else: + self.model_runner = GPUARModelRunner(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + @instrument(span_name="Compile/warmup") + def compile_or_warm_up_model(self) -> float: + """Skip warmup_kernels for V2 Omni AR models. + + Upstream ``compile_or_warm_up_model`` calls ``warmup_kernels()`` + only for V2 model runners. ``warmup_kernels`` creates fake + requests with hardcoded block_ids that bypass KVCacheManager + and runs real (non-dummy) execute_model + sample_tokens. For + Omni AR models with preprocess, this pollutes model-internal + state and produces incorrect logits for subsequent real + requests (observed 30x more decode steps than V1). + + This override runs the parent implementation but patches out + the ``warmup_kernels`` call, matching V1 behavior. + """ + if not self.use_v2_model_runner: + return super().compile_or_warm_up_model() + + # Temporarily disable warmup_kernels during parent call. + # Must patch both the module attribute AND the local name in + # gpu_worker (from-import creates a local binding). + import vllm.v1.worker.gpu_worker as _gw + + _orig = _gw.warmup_kernels + + def _noop(*a, **kw): + return None + + _gw.warmup_kernels = _noop + # Also patch the module where it was originally defined + import vllm.v1.worker.gpu.warmup as _wm + + _orig_wm = _wm.warmup_kernels + _wm.warmup_kernels = _noop + try: + result = super().compile_or_warm_up_model() + finally: + _gw.warmup_kernels = _orig + _wm.warmup_kernels = _orig_wm + logger.info("compile_or_warm_up_model: skipped warmup_kernels for Omni AR V2") + return result + def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK: """ Explicitly handle sleep commands. diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 2a85436a0..60aeea458 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -11,6 +11,7 @@ from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.compat import make_filtered_namedtuple from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -18,6 +19,21 @@ from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) +VLLM_OMNI_USE_V2_RUNNER = bool(int(os.environ.get("VLLM_OMNI_USE_V2_RUNNER", "0"))) + + +def _make_compilation_times(language_model_time: float, **kwargs): + try: + from vllm.v1.worker.gpu_worker import CompilationTimes + except ImportError: + return language_model_time + + values = {"language_model": language_model_time, "encoder": 0.0, **kwargs} + result, unknown = make_filtered_namedtuple(CompilationTimes, **values) + if unknown: + logger.warning("Unknown fields passed to CompilationTimes: %s", sorted(unknown)) + return result + class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU Worker for Generation model (non-autoregressive waveform generation). @@ -28,7 +44,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): @instrument(span_name="Init device") def init_device(self): - if self.device_config.device_type in ("cuda", "musa"): + if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) parallel_config = self.parallel_config @@ -51,7 +67,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): assert self.local_rank < torch.accelerator.device_count(), ( f"DP adjusted local rank {self.local_rank} is out of bounds. " ) - visible_device_count = torch.accelerator.device_count() + visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0 assert self.parallel_config.local_world_size <= visible_device_count, ( f"local_world_size ({self.parallel_config.local_world_size}) must " f"be less than or equal to the number of visible devices " @@ -93,13 +109,28 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) - if self.use_v2_model_runner: - # OMNI: v2 model runner does not yet include omni hooks. - logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False + if VLLM_OMNI_USE_V2_RUNNER or self.use_v2_model_runner: + from vllm_omni.worker_v2.omni_generation_model_runner import ( + OmniGenerationModelRunner, + ) - self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) + logger.info("Using MR v2 OmniGenerationModelRunner for generation stage.") + self.use_v2_model_runner = True + self.model_runner = OmniGenerationModelRunner(self.vllm_config, self.device) + else: + self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + + @instrument(span_name="Compile/warmup") + def compile_or_warm_up_model(self) -> float: + """Generation stages have no KV cache or sampler — skip warmup_kernels.""" + if not self.use_v2_model_runner: + return super().compile_or_warm_up_model() + import time + + start = time.perf_counter() + self.model_runner.profile_run() + return _make_compilation_times(time.perf_counter() - start) diff --git a/vllm_omni/worker_v2/__init__.py b/vllm_omni/worker_v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker_v2/forward_compat.py b/vllm_omni/worker_v2/forward_compat.py new file mode 100644 index 000000000..c9deecbb9 --- /dev/null +++ b/vllm_omni/worker_v2/forward_compat.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Forward-call compatibility helpers for Omni MR V2 runners.""" + +from __future__ import annotations + +from typing import Any + + +def add_forward_compat_kwargs( + model_inputs: dict[str, Any], + input_batch: Any, + sampler: Any, +) -> None: + """Add Omni model forward kwargs that V1 model runners provided. + + Some Omni model stages inspect these kwargs inside ``forward`` rather + than only in the runner. Qwen2.5-Omni's talker, for example, rewrites + text-model pad ids in ``sampling_metadata`` to a talker-vocab pad id + before sampling codec tokens. + """ + model_inputs.setdefault( + "sampling_metadata", + getattr(input_batch, "sampling_metadata", None), + ) + logits_index = getattr( + input_batch, + "logits_indices", + getattr(input_batch, "logits_index", None), + ) + model_inputs.setdefault("logits_index", logits_index) + model_inputs.setdefault("sampler", sampler) diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py new file mode 100644 index 000000000..79baf4f6a --- /dev/null +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -0,0 +1,49 @@ +"""Omni-aware ``init_model_state`` factory. + +Extends the upstream v2 factory with Omni architecture dispatch. +Non-Omni architectures fall through to the upstream ``init_model_state``. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states import ( + init_model_state as _upstream_init_model_state, +) +from vllm.v1.worker.gpu.model_states.interface import ModelState + +# Keep in sync: when adding a new Omni model architecture, add it here +# AND update the corresponding test in tests/worker_v2/test_init_model_state.py. +_OMNI_ARCHITECTURES: set[str] = { + "Qwen3OmniMoeForConditionalGeneration", + "Qwen2_5OmniForConditionalGeneration", + "MammothModa2ForConditionalGeneration", + "MiMoAudioForConditionalGeneration", + "MammothModa2ARForConditionalGeneration", + "Qwen3TTSTalkerForConditionalGeneration", + "Qwen3TTSCode2Wav", +} + + +def init_omni_model_state( + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, +) -> ModelState: + """Create the appropriate ``ModelState`` for *model*. + + Returns an ``OmniModelState`` when the configured architecture is a + known Omni model; otherwise delegates to the upstream v2 factory. + """ + archs = set(vllm_config.model_config.architectures or []) + if archs & _OMNI_ARCHITECTURES: + from vllm_omni.worker_v2.model_states.omni_model_state import ( + OmniModelState, + ) + + return OmniModelState(vllm_config, model, encoder_cache, device) + return _upstream_init_model_state(vllm_config, model, encoder_cache, device) diff --git a/vllm_omni/worker_v2/model_states/intermediate_buffer.py b/vllm_omni/worker_v2/model_states/intermediate_buffer.py new file mode 100644 index 000000000..1e8743c78 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/intermediate_buffer.py @@ -0,0 +1,147 @@ +"""OmniIntermediateBuffer — per-request cross-stage state for Omni pipelines. + +Uses ``req_index`` (not ``req_id``) for O(1) access, aligned with v2's +``RequestState`` slot management. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.input_batch import InputBatch + +logger = init_logger(__name__) + + +def _resolve_prompt_embeds(pe: Any) -> torch.Tensor | None: + """Convert a prompt_embeds payload to a contiguous CPU tensor.""" + if pe is None: + return None + try: + if isinstance(pe, torch.Tensor): + return pe.detach().cpu().contiguous() + data = getattr(pe, "data", None) + shape = getattr(pe, "shape", None) + if data is not None and shape is not None: + dt = np.dtype(getattr(pe, "dtype", "float32")) + arr = np.frombuffer(data, dtype=dt).reshape(shape) + return torch.from_numpy(arr.copy()) + except Exception: + logger.exception("Failed to decode prompt_embeds payload") + return None + + +def _resolve_additional_information(payload: Any) -> dict[str, Any]: + """Convert an additional_information payload to a plain dict.""" + if payload is None: + return {} + try: + from vllm_omni.data_entry_keys import unflatten_payload + from vllm_omni.engine.serialization import deserialize_additional_information + + info = deserialize_additional_information(payload) + if isinstance(info, dict) and any(isinstance(key, str) and "." in key for key in info): + info = unflatten_payload(info) + return info if isinstance(info, dict) else {} + except Exception: + logger.exception("Failed to decode additional_information payload") + return {} + + +class OmniIntermediateBuffer: + """Per-request intermediate state for multi-stage Omni pipelines. + + Stores ``prompt_embeds``, ``additional_information``, ``mm_features``, + ``req_id`` and any runtime updates written back by model postprocess. + """ + + def __init__(self, max_num_reqs: int): + self.buffers: list[dict[str, Any]] = [{} for _ in range(max_num_reqs)] + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + info: dict[str, Any] = {} + + pe = getattr(new_req_data, "prompt_embeds", None) + if pe is not None: + pe_cpu = _resolve_prompt_embeds(pe) + if pe_cpu is not None: + info["prompt_embeds_cpu"] = pe_cpu + + ai = getattr(new_req_data, "additional_information", None) + if ai is not None: + info.update(_resolve_additional_information(ai)) + + if new_req_data.mm_features: + info["mm_features"] = new_req_data.mm_features + + info["req_id"] = new_req_data.req_id + self.buffers[req_index] = info + + def remove_request(self, req_index: int) -> None: + self.buffers[req_index] = {} + + def gather(self, input_batch: InputBatch) -> list[dict[str, Any]]: + """Return buffer dicts in current batch order (via ``idx_mapping_np``).""" + return [self.buffers[idx] for idx in input_batch.idx_mapping_np] + + @staticmethod + def _split_gpu_keys(gpu_resident_keys: set[Any] | None) -> tuple[set[Any], dict[Any, set[Any]]]: + top_level: set[Any] = set() + nested: dict[Any, set[Any]] = {} + for key in gpu_resident_keys or set(): + if isinstance(key, tuple) and len(key) == 2: + nested.setdefault(key[0], set()).add(key[1]) + else: + top_level.add(key) + return top_level, nested + + @staticmethod + def _store_value(dest: dict[Any, Any], key: Any, value: Any, gpu_keys: set[Any]) -> None: + if isinstance(value, torch.Tensor): + if key in gpu_keys: + dest[key] = value.detach().clone() + else: + dest[key] = value.detach().cpu().contiguous() + elif isinstance(value, list): + dest[key] = [ + (item.detach().cpu().contiguous() if isinstance(item, torch.Tensor) else item) for item in value + ] + else: + dest[key] = value + + def update( + self, + req_index: int, + updates: dict[Any, Any], + gpu_resident_keys: set[Any] | None = None, + ) -> None: + """Merge *updates* into the buffer at *req_index*. + + Tensors are detached; those whose key is **not** in + *gpu_resident_keys* are moved to CPU. + """ + if not updates: + return + top_gpu_keys, nested_gpu_keys = self._split_gpu_keys(gpu_resident_keys) + existing = self.buffers[req_index] + for k, v in updates.items(): + if isinstance(k, tuple) and len(k) == 2: + type_key, qualifier = k + existing_sub = existing.setdefault(type_key, {}) + if not isinstance(existing_sub, dict): + existing_sub = {} + existing[type_key] = existing_sub + self._store_value(existing_sub, qualifier, v, nested_gpu_keys.get(type_key, set())) + elif isinstance(v, dict): + existing_sub = existing.setdefault(k, {}) + if not isinstance(existing_sub, dict): + existing_sub = {} + existing[k] = existing_sub + for qualifier, value in v.items(): + self._store_value(existing_sub, qualifier, value, nested_gpu_keys.get(k, set())) + else: + self._store_value(existing, k, v, top_gpu_keys) diff --git a/vllm_omni/worker_v2/model_states/models/__init__.py b/vllm_omni/worker_v2/model_states/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py new file mode 100644 index 000000000..ccb4937dd --- /dev/null +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -0,0 +1,513 @@ +"""OmniModelState — generic ModelState base for all Omni model stages. + +Extends ``DefaultModelState`` with: + +* Cross-stage intermediate buffer (``OmniIntermediateBuffer``) +* ``model_intermediate_buffer`` / ``runtime_additional_information`` injection + into ``model_inputs`` via ``prepare_inputs()`` +* ``OmniOutput`` → ``(text_hidden, multimodal_outputs)`` post-processing +* Plugin lifecycle dispatch (``OmniModelStatePlugin``) +""" + +from __future__ import annotations + +import threading +import types +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from vllm.v1.worker.utils import AttentionGroup + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.states import RequestState + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + OmniIntermediateBuffer, +) +from vllm_omni.worker_v2.model_states.plugin import OmniModelStatePlugin + +logger = init_logger(__name__) +_rope_patch_lock = threading.Lock() + + +def _default_mrope_positions( + self_model: Any, + input_tokens: list[int], + mm_features: list, +) -> tuple[torch.Tensor, int]: + """Return 3D sequential positions with zero delta. + + For non-vision Omni models (e.g. TTS Talker), all 3 M-RoPE + dimensions use the same sequential positions. Delta=0 keeps decode + positions sequential, identical to the 1D case but broadcast to 3 dims. + """ + n = len(input_tokens) + pos = torch.arange(n, dtype=torch.long) + return pos.unsqueeze(0).expand(3, -1), 0 + + +def _make_safe_get_rope(orig_get_rope): + from vllm.v1.worker.gpu.mm.rope import RopeState + + def _safe_get_rope(model_config: Any, mdl: Any, **kwargs: Any) -> Any: + try: + result = orig_get_rope(model_config, mdl, **kwargs) + except (AssertionError, TypeError): + result = None + + needs_mrope = bool(getattr(model_config, "uses_mrope", False)) + if result is not None and (not needs_mrope or getattr(result, "num_dims", 0) >= 3): + return result + if not needs_mrope: + return None + if not hasattr(mdl, "get_mrope_input_positions"): + mdl.get_mrope_input_positions = types.MethodType(_default_mrope_positions, mdl) + return RopeState(num_dims=3, has_delta=True, **kwargs) + + return _safe_get_rope + + +class OmniModelState(DefaultModelState): + """Generic Omni ``ModelState`` — works for **all** Omni model stages. + + Model-specific behaviour is injected via ``OmniModelStatePlugin`` + instances or subclasses; this class itself is model-agnostic. + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + # DefaultModelState.__init__ calls get_rope_state() which asserts + # isinstance(model, SupportsMRoPE). Two categories of Omni models: + # + # 1. Models that implement SupportsMRoPE (e.g. Qwen3-Omni Thinker): + # get_rope_state() succeeds normally, _safe_get_rope is a no-op. + # These models get correct 3D M-RoPE positions from the runner. + # + # 2. Models that do NOT implement SupportsMRoPE (e.g. Qwen3-TTS + # Talker, Code2Wav, FishSpeech): get_rope_state() would assert. + # These models compute their own position encoding internally + # (via model.forward kwargs or fixed 1D positions from + # InputBatch.positions), so rope_state = None is correct — + # DefaultModelState.prepare_inputs returns {} when rope_state + # is None, and upstream execute_model falls back to + # InputBatch.positions (1D sequential). + # Patch get_rope_state to handle Omni models that declare + # M-RoPE in config (mrope_section) but do not implement the + # SupportsMRoPE interface. For these models we create a + # RopeState with 3D sequential positions (matching V1 MR). + # + # The patch is applied via a class-level lock to prevent + # concurrent OmniModelState instances (e.g. different stages + # in a thread pool) from overwriting each other's patch. + from vllm.v1.worker.gpu.model_states import default as _default_mod + + with _rope_patch_lock: + orig_get_rope = _default_mod.get_rope_state + _default_mod.get_rope_state = _make_safe_get_rope(orig_get_rope) + try: + super().__init__(vllm_config, model, encoder_cache, device) + finally: + _default_mod.get_rope_state = orig_get_rope + max_num_reqs = self.scheduler_config.max_num_seqs + self.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) + self.has_preprocess: bool = getattr(model, "has_preprocess", False) + self.has_postprocess: bool = getattr(model, "has_postprocess", False) + self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False) + self.plugins: list[OmniModelStatePlugin] = [] + + # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim. + self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0 + + # Static inputs_embeds buffer for FULL CUDA graph — preprocess fills it in-place each step. + self._static_inputs_embeds: torch.Tensor | None = None + if self._embed_dim > 0: + self._static_inputs_embeds = torch.zeros( + (self.max_num_tokens, self._embed_dim), + dtype=self.dtype, + device=device, + ) + + # Static MTP buffers so _run_batched_mtp uses .copy_() instead of torch.cat(). + self._mtp_input_ids: torch.Tensor | None = None + self._mtp_input_embeds: torch.Tensor | None = None + self._mtp_hidden: torch.Tensor | None = None + self._mtp_text_step: torch.Tensor | None = None + if self._embed_dim > 0 and hasattr(model, "talker_mtp"): + max_bs = max_num_reqs + self._mtp_input_ids = torch.zeros(max_bs, dtype=torch.long, device=device) + self._mtp_input_embeds = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + self._mtp_hidden = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + self._mtp_text_step = torch.zeros((max_bs, self._embed_dim), dtype=self.dtype, device=device) + + if hasattr(model, "get_omni_plugins"): + for plugin in model.get_omni_plugins(): + self.register_plugin(plugin) + + @staticmethod + def _get_embed_dim(model: nn.Module, device: torch.device) -> int: + """Return the embedding dim that ``embed_input_ids`` produces (may differ from hf_text_config).""" + if hasattr(model, "embed_input_ids"): + dummy = torch.zeros(1, dtype=torch.long, device=device) + with torch.no_grad(): + out = model.embed_input_ids(dummy) + return out.shape[-1] + return 0 + + # ------------------------------------------------------------------ + # Attention metadata: use actual max_seq_len, not max_model_len + # ------------------------------------------------------------------ + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: Any, + for_capture: bool = False, + ) -> dict[str, Any]: + """Override to use actual max_seq_len instead of max_model_len. + + Upstream DefaultModelState uses ``self.max_model_len`` as + ``max_seq_len`` for attention metadata. This causes + FlashAttention to use a different scheduler/tiling strategy + than V1 MR (which uses the actual maximum sequence length). + The different tiling changes the float accumulation order in + bf16, causing numerically different attention outputs that + snowball through the transformer layers and produce completely + different logits—leading to 30x more decode steps for TTS. + """ + from vllm.v1.worker.gpu.attn_utils import build_attn_metadata + + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + + # Use actual max seq_len (matching V1 MR behavior) instead of + # max_model_len. For CUDA graph capture, use max_model_len to + # ensure the captured kernel covers all possible seq_lens. + if for_capture: + max_seq_len = self.max_model_len + else: + max_seq_len = int(input_batch.seq_lens[:num_reqs].max().item()) + + return build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=max_seq_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + ) + + # ------------------------------------------------------------------ + # Plugin management + # ------------------------------------------------------------------ + + def register_plugin(self, plugin: OmniModelStatePlugin) -> None: + self.plugins.append(plugin) + + # ------------------------------------------------------------------ + # Request lifecycle + # ------------------------------------------------------------------ + + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: + super().add_request(req_index, new_req_data) + self.intermediate_buffer.add_request(req_index, new_req_data) + for plugin in self.plugins: + plugin.on_add_request(req_index, new_req_data) + + def remove_request(self, req_index: int) -> None: + self.intermediate_buffer.remove_request(req_index) + for plugin in self.plugins: + plugin.on_remove_request(req_index) + + # ------------------------------------------------------------------ + # Input preparation + # ------------------------------------------------------------------ + + def prepare_inputs(self, input_batch: InputBatch, req_states: RequestState) -> dict[str, Any]: + base = super().prepare_inputs(input_batch, req_states) + buffer_list = self.intermediate_buffer.gather(input_batch) + base["model_intermediate_buffer"] = buffer_list + base["runtime_additional_information"] = buffer_list + base["seq_token_counts"] = [int(input_batch.num_scheduled_tokens[i]) for i in range(input_batch.num_reqs)] + # Return static inputs_embeds so FULL graph replay uses the same + # tensor address that was captured. Preprocess fills it in-place. + if self._static_inputs_embeds is not None: + base["inputs_embeds"] = self._static_inputs_embeds[: input_batch.num_tokens_after_padding] + for plugin in self.plugins: + base.update(plugin.prepare_extra_inputs(input_batch, req_states)) + return base + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + base = super().prepare_dummy_inputs(num_reqs, num_tokens) + dummy_buffer = [{} for _ in range(num_reqs)] + base["model_intermediate_buffer"] = dummy_buffer + base["runtime_additional_information"] = dummy_buffer + if num_reqs > 0: + per_req = num_tokens // num_reqs + remainder = num_tokens % num_reqs + counts = [per_req] * num_reqs + counts[-1] += remainder + base["seq_token_counts"] = counts + else: + base["seq_token_counts"] = [] + # Return static inputs_embeds for FULL graph capture so the graph + # captures this tensor's address. + if self._static_inputs_embeds is not None: + base["inputs_embeds"] = self._static_inputs_embeds[:num_tokens] + return base + + # ------------------------------------------------------------------ + # Pre-forward: per-request preprocess + batched MTP + # ------------------------------------------------------------------ + + def run_preprocess(self, input_batch: InputBatch, model_inputs: dict[str, Any]) -> None: + """Per-request preprocess + MTP before model forward. + + Modifies ``model_inputs["input_ids"]`` and ``model_inputs["inputs_embeds"]`` + in-place. Collects decode-step MTP inputs and runs a single batched MTP + forward at the end. + + Skipped when the model declares ``preprocess_in_forward = True``, + meaning it handles preprocess internally inside forward(). + """ + if not self.has_preprocess: + return + # Model does preprocess+MTP inside forward() — skip external preprocess. + if getattr(self.model, "preprocess_in_forward", False): + return + + embeds = model_inputs.get("inputs_embeds") + if embeds is None: + embeds = self.model.embed_input_ids(input_batch.input_ids[: input_batch.num_tokens]) + model_inputs["inputs_embeds"] = embeds + + input_ids = model_inputs.get("input_ids") + if input_ids is None: + input_ids = input_batch.input_ids + + gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set()) + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]] = [] + + for i in range(input_batch.num_reqs): + req_idx = int(input_batch.idx_mapping_np[i]) + buf = self.intermediate_buffer.buffers[req_idx] + + # Skip warmup/dummy requests that have no real buffer data. + # Warmup creates fake requests without the metadata that + # model.preprocess() requires (e.g. additional_information.text). + if not buf or "req_id" not in buf: + continue + + start = int(input_batch.query_start_loc_np[i]) + n_tok = int(input_batch.num_scheduled_tokens[i]) + + ids_slice = input_ids[start : start + n_tok] + emb_slice = embeds[start : start + n_tok] + + try: + info = {key: value for key, value in buf.items() if isinstance(key, str)} + new_ids, new_emb, updates = self.model.preprocess(ids_slice, emb_slice, **info) + except Exception: + logger.warning( + "preprocess failed for req_idx=%d (req_id=%s); skipping preprocess for this request", + req_idx, + buf.get("req_id", "?"), + exc_info=True, + ) + continue + + # Write back in-place + seg = min(n_tok, new_ids.shape[0]) + input_ids[start : start + seg] = new_ids[:seg] + embeds[start : start + seg] = new_emb[:seg] + + # Collect MTP inputs for decode steps (n_tok == 1 with mtp_inputs) + mtp_inputs = updates.pop("mtp_inputs", None) + if mtp_inputs is not None and n_tok == 1: + mtp_batches.append((i, start, mtp_inputs)) + + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + if mtp_batches and hasattr(self.model, "talker_mtp"): + self._run_batched_mtp(mtp_batches, input_ids, embeds, input_batch, gpu_keys) + + def _run_batched_mtp( + self, + mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], + input_ids: torch.Tensor, + embeds: torch.Tensor, + input_batch: InputBatch, + gpu_keys: set[str], + ) -> None: + """Batch MTP forward for all decode-step requests. + + Uses pre-allocated static buffers to avoid per-step torch.cat + memory allocations. + """ + from vllm.forward_context import set_forward_context + + bsz = len(mtp_batches) + + # Fill static buffers via .copy_() — no allocation. + if self._mtp_input_ids is not None and bsz <= self._mtp_input_ids.shape[0]: + for j, (_i, start, (past_hidden, text_step)) in enumerate(mtp_batches): + self._mtp_input_ids[j] = input_ids[start] + self._mtp_input_embeds[j].copy_(embeds[start]) + self._mtp_hidden[j].copy_(past_hidden.reshape(-1)[: self._mtp_hidden.shape[1]]) + self._mtp_text_step[j].copy_(text_step.reshape(-1)[: self._mtp_text_step.shape[1]]) + batch_ids = self._mtp_input_ids[:bsz] + batch_emb = self._mtp_input_embeds[:bsz] + batch_hidden = self._mtp_hidden[:bsz] + batch_step = self._mtp_text_step[:bsz] + else: + # Fallback to torch.cat if static buffers not available. + ids_list, emb_list, hidden_list, step_list = [], [], [], [] + for _i, start, (past_hidden, text_step) in mtp_batches: + ids_list.append(input_ids[start : start + 1]) + emb_list.append(embeds[start : start + 1]) + hidden_list.append(past_hidden) + step_list.append(text_step) + batch_ids = torch.cat(ids_list) + batch_emb = torch.cat(emb_list) + batch_hidden = torch.cat(hidden_list) + batch_step = torch.cat(step_list) + + with set_forward_context( + None, + self.vllm_config, + num_tokens=bsz, + ): + new_emb, codes = self.model.talker_mtp( + batch_ids, + batch_emb, + batch_hidden, + batch_step, + ) + + audio_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio")) + for j, (i, start, _) in enumerate(mtp_batches): + embeds[start : start + 1] = new_emb[j : j + 1] + req_idx = int(input_batch.idx_mapping_np[i]) + if isinstance(audio_key, tuple) and len(audio_key) == 2: + updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}} + elif isinstance(audio_key, str): + updates = {audio_key: codes[j : j + 1]} + else: + raise TypeError( + f"talker_mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}" + ) + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + # ------------------------------------------------------------------ + # Post-forward: per-request postprocess + # ------------------------------------------------------------------ + + def run_postprocess(self, hidden_states: torch.Tensor, input_batch: InputBatch) -> None: + """Per-request postprocess after model forward. + + Extracts per-request updates from hidden_states and writes them + back to the intermediate buffer (e.g. ``last_talker_hidden``). + + Skipped when the model declares ``preprocess_in_forward = True`` + (the flag covers both pre- and post-processing — both run inside + the model's forward()). + """ + if not self.has_postprocess: + return + # preprocess_in_forward also covers postprocess — both run inside forward() + if getattr(self.model, "preprocess_in_forward", False): + return + gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set()) + for i in range(input_batch.num_reqs): + req_idx = int(input_batch.idx_mapping_np[i]) + buf = self.intermediate_buffer.buffers[req_idx] + if not buf or "req_id" not in buf: + continue + start = int(input_batch.query_start_loc_np[i]) + n_tok = int(input_batch.num_scheduled_tokens[i]) + h_slice = hidden_states[start : start + n_tok] + info = {key: value for key, value in buf.items() if isinstance(key, str) and key != "hidden_states"} + updates = self.model.postprocess(h_slice, **info) + if updates: + self.intermediate_buffer.update(req_idx, updates, gpu_keys) + + # ------------------------------------------------------------------ + # Output post-processing + # ------------------------------------------------------------------ + + def postprocess_model_output( + self, + model_output: Any, + input_batch: InputBatch, + req_states: RequestState, + ) -> tuple[torch.Tensor, dict]: + """Convert raw model output to ``(text_hidden, multimodal_outputs)``. + + Handles ``OmniOutput`` unwrapping and ``make_omni_output`` + conversion, then dispatches to registered plugins. + """ + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + if isinstance(model_output, (list, tuple)) or self.have_multimodal_outputs: + buffer_list = self.intermediate_buffer.gather(input_batch) + try: + model_output = self.model.make_omni_output( + model_output, + model_intermediate_buffer=buffer_list, + runtime_additional_information=buffer_list, + ) + except Exception: + _desc = type(model_output).__name__ + if isinstance(model_output, (list, tuple)): + _desc += f"(len={len(model_output)})" + logger.warning( + "make_omni_output failed for %s; multimodal outputs will be empty", + _desc, + exc_info=True, + ) + + if isinstance(model_output, OmniOutput): + text_hidden = model_output.text_hidden_states + multimodal_outputs: dict = model_output.multimodal_outputs or {} + elif isinstance(model_output, (list, tuple)): + text_hidden = model_output[0] + multimodal_outputs = {} + else: + text_hidden = model_output + multimodal_outputs = {} + + for plugin in self.plugins: + text_hidden, multimodal_outputs = plugin.postprocess( + text_hidden, multimodal_outputs, input_batch, req_states + ) + + return text_hidden, multimodal_outputs diff --git a/vllm_omni/worker_v2/model_states/plugin.py b/vllm_omni/worker_v2/model_states/plugin.py new file mode 100644 index 000000000..6c5d90112 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/plugin.py @@ -0,0 +1,42 @@ +"""OmniModelStatePlugin — abstract extension point for model-specific behavior. + +Plugins are registered on OmniModelState at model-load time and receive +callbacks at well-defined points in the request/step lifecycle. The base +class provides default no-op implementations for optional hooks so that +concrete plugins only need to override what they use. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.states import RequestState + + +class OmniModelStatePlugin(ABC): + """Extension point for model-specific behavior in OmniModelState.""" + + @abstractmethod + def on_add_request(self, req_index: int, new_req_data: NewRequestData) -> None: ... + + @abstractmethod + def on_remove_request(self, req_index: int) -> None: ... + + def prepare_extra_inputs(self, input_batch: InputBatch, req_states: RequestState) -> dict[str, Any]: + return {} + + def postprocess( + self, + text_hidden: Any, + multimodal_outputs: dict, + input_batch: InputBatch, + req_states: RequestState, + ) -> tuple[Any, dict]: + """Post-process model output. Default is pass-through.""" + return text_hidden, multimodal_outputs + + def dummy_run(self, state: Any, num_tokens: int) -> None: + pass diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py new file mode 100644 index 000000000..31af092f5 --- /dev/null +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -0,0 +1,564 @@ +"""OmniARModelRunner — autoregressive stage runner on MR V2. + +Extends ``OmniGPUModelRunner`` with: + +* ``OmniOutput`` post-processing in ``sample_tokens`` +* Per-request ``pooler_output`` construction (hidden + multimodal slices) +* Async D2H copy via ``OmniAsyncOutput`` for non-blocking output transfer +* Cross-stage KV extraction before state cleanup +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput + +from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVTransferManager, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +logger = init_logger(__name__) + + +class OmniARModelRunner(OmniGPUModelRunner): + """AR stage runner. Produces per-request hidden states + multimodal outputs.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.kv_transfer_manager: OmniKVTransferManager | None = None + self._kv_extracted_req_ids: list[str] | None = None + + def _ensure_kv_transfer_manager(self) -> OmniKVTransferManager: + if self.kv_transfer_manager is None: + self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + return self.kv_transfer_manager + + # ------------------------------------------------------------------ + # execute_model: KV transfer pre-hook + delegate to super + # ------------------------------------------------------------------ + + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: Any | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + ) -> Any: + if not dummy_run: + self._handle_kv_transfer_pre(scheduler_output) + return super().execute_model( + scheduler_output, + intermediate_tensors, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + is_profile=is_profile, + ) + + # ------------------------------------------------------------------ + # sample_tokens: OmniOutput handling + pooler_output + async D2H + # ------------------------------------------------------------------ + + def sample_tokens( + self, grammar_output: GrammarOutput | None + ) -> OmniAsyncOutput | OmniModelRunnerOutput | ModelRunnerOutput | None: + kv_extracted = self._kv_extracted_req_ids + self._kv_extracted_req_ids = None + + if self.execute_model_state is None: + return None + + input_batch = self.execute_model_state.input_batch + hidden_states = self.execute_model_state.hidden_states + # kv_connector_output is no longer a field on vLLM 0.23.0's ExecuteModelState; + # the base execute_model stashes it here after post_forward(). + kv_connector_output = self._last_kv_connector_output + self._last_kv_connector_output = None + self.execute_model_state = None + + if not self.is_last_pp_rank: + from vllm.v1.worker.gpu.pp_utils import pp_receive + + sampled, num_sampled, num_rejected = pp_receive( + input_batch.num_reqs, + max_sample_len=self.num_speculative_steps + 1, + ) + # vLLM 0.23.0 renamed postprocess() -> postprocess_sampled() and now + # takes idx_mapping (+ query_start_loc) instead of the InputBatch. + self.postprocess_sampled( + input_batch.idx_mapping, + sampled, + num_sampled, + num_rejected, + input_batch.query_start_loc, + ) + return None + + # --- Omni: reconstruct raw model output and post-process --- + aux = self._last_aux_output + self._last_aux_output = None + multimodal_outputs = self._last_multimodal_outputs + self._last_multimodal_outputs = None + raw_output = self._reconstruct_raw_model_output( + hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + aux=aux, + ) + text_hidden, multimodal_outputs = self.model_state.postprocess_model_output( + raw_output, input_batch, self.req_states + ) + + # --- Standard v2 sampling --- + sampler_output, num_sampled, num_rejected = self._sample_with_prompt_token_compat( + text_hidden, + input_batch, + grammar_output, + ) + + if self.use_pp: + from vllm.v1.worker.gpu.pp_utils import pp_broadcast + + pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected) + + # --- Omni: prompt logprobs --- + assert self.prompt_logprobs_worker is not None + # vLLM 0.23.0 dropped the prefill_len / num_computed_prefill_tokens args; + # compute_prompt_logprobs now reads prefill_len_np / + # num_computed_prefill_tokens_np from input_batch instead. Mirror the + # parent GPUModelRunner's 6-arg call exactly. + prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs( + self.model.compute_logits, + text_hidden, + input_batch, + self.req_states.all_token_ids.gpu, + self.req_states.num_computed_tokens.gpu, + self.req_states.prompt_len.np, + ) + + # --- Omni: pooler_output --- + engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", "text") + need_pooler = engine_output_type != "text" + + # --- Build base output --- + model_runner_output = OmniModelRunnerOutput( + req_ids=input_batch.req_ids, + req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)}, + sampled_token_ids=None, # type: ignore[arg-type] + prompt_logprobs_dict=prompt_logprobs_dict, + kv_connector_output=kv_connector_output, + ) + model_runner_output.kv_extracted_req_ids = kv_extracted + + # --- Async D2H via OmniAsyncOutput --- + async_output = OmniAsyncOutput( + model_runner_output=model_runner_output, + sampler_output=sampler_output, + num_sampled_tokens=num_sampled, + main_stream=self.main_stream, + copy_stream=self.output_copy_stream, + text_hidden=text_hidden if need_pooler else None, + multimodal_outputs=multimodal_outputs if need_pooler else None, + input_batch=input_batch if need_pooler else None, + ) + + # Postprocess AFTER creating async output (so copy_event is + # recorded before postprocess, matching upstream pattern). + self.postprocess_sampled( + input_batch.idx_mapping, + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch.query_start_loc, + ) + + if self.use_async_scheduling: + return async_output + return async_output.get_output() + + # ------------------------------------------------------------------ + # pooler_output construction + # ------------------------------------------------------------------ + + @staticmethod + def _reconstruct_raw_model_output( + *, + hidden_states: torch.Tensor, + multimodal_outputs: dict[str, Any] | None, + aux: Any | None, + ) -> Any: + if multimodal_outputs: + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + ) + if aux is not None: + return (hidden_states, aux) + return hidden_states + + @staticmethod + def _build_pooler_output_from_cpu( + hidden_cpu: torch.Tensor, + mm_cpu: dict[str, Any], + query_start_loc_np: np.ndarray, + num_scheduled_tokens: np.ndarray, + num_reqs: int, + ) -> list[dict[str, Any]]: + """Build pooler_output from already-CPU tensors.""" + total = hidden_cpu.shape[0] + pooler: list[dict[str, Any]] = [] + for i in range(num_reqs): + start = int(query_start_loc_np[i]) + end = start + int(num_scheduled_tokens[i]) + payload: dict[str, Any] = {"hidden": hidden_cpu[start:end]} + for k, v in mm_cpu.items(): + payload[k] = _slice_pooler_value( + v, + req_index=i, + start=start, + end=end, + total_tokens=total, + ) + pooler.append(flatten_payload(payload)) + return pooler + + def _sample_with_prompt_token_compat( + self, + text_hidden: torch.Tensor, + input_batch: Any, + grammar_output: GrammarOutput | None, + ) -> tuple[Any, Any, Any]: + """Run upstream sampling while restoring V1 prompt-id compatibility. + + Some Omni AR stages sample from a logits vocabulary smaller than the + tokenizer/input vocabulary. V1 corrected ``prompt_token_ids`` after it + had the real logits tensor. MR V2's upstream ``sample`` owns logits + computation, so wrap ``compute_logits`` for this call and clamp from the + actual logits shape instead of guessing from config. + """ + compute_logits = getattr(self.model, "compute_logits", None) + if not callable(compute_logits): + return self.sample(text_hidden, input_batch, grammar_output) + + def compute_logits_with_prompt_token_compat(*args: Any, **kwargs: Any) -> Any: + logits = compute_logits(*args, **kwargs) + logits_shape = getattr(logits, "shape", ()) if logits is not None else () + logits_vocab_size = logits_shape[-1] if logits_shape else None + if isinstance(logits_vocab_size, int): + self._clamp_sampling_prompt_token_ids(input_batch, logits_vocab_size) + return logits + + model_dict = getattr(self.model, "__dict__", {}) + had_instance_compute_logits = isinstance(model_dict, dict) and "compute_logits" in model_dict + original_instance_compute_logits = model_dict.get("compute_logits") if had_instance_compute_logits else None + setattr(self.model, "compute_logits", compute_logits_with_prompt_token_compat) + try: + return self.sample(text_hidden, input_batch, grammar_output) + finally: + if had_instance_compute_logits: + setattr(self.model, "compute_logits", original_instance_compute_logits) + else: + try: + delattr(self.model, "compute_logits") + except AttributeError: + setattr(self.model, "compute_logits", compute_logits) + + @staticmethod + def _clamp_sampling_prompt_token_ids( + input_batch: Any, + logits_vocab_size: int | None, + ) -> None: + """Clamp sampler prompt IDs to the stage logits vocabulary. + + V1 Omni AR runner did this after computing logits. MR V2 samples via + the upstream helper, so normalize the metadata before that call. + """ + if logits_vocab_size is None or logits_vocab_size <= 0: + return + if getattr(input_batch, "vocab_size", logits_vocab_size) <= logits_vocab_size: + return + + sampling_metadata = getattr(input_batch, "sampling_metadata", None) + if sampling_metadata is None or getattr(sampling_metadata, "no_penalties", False): + return + + prompt_token_ids = getattr(sampling_metadata, "prompt_token_ids", None) + if prompt_token_ids is None: + return + + max_token_id = logits_vocab_size - 1 + if isinstance(prompt_token_ids, torch.Tensor): + prompt_token_ids.clamp_(max=max_token_id) + return + + if isinstance(prompt_token_ids, list): + sampling_metadata.prompt_token_ids = [ + [min(int(tok), max_token_id) for tok in ids] if isinstance(ids, list) else min(int(ids), max_token_id) + for ids in prompt_token_ids + ] + + # ------------------------------------------------------------------ + # KV transfer + # ------------------------------------------------------------------ + + def _handle_kv_transfer_pre(self, scheduler_output: SchedulerOutput) -> None: + finished: dict = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) + if not finished: + return + + kv_caches = getattr(self, "kv_caches", None) + if kv_caches is None: + return + + if hasattr(self.model, "get_kv_transfer_metadata"): + for req_id, data in finished.items(): + try: + meta = self.model.get_kv_transfer_metadata(req_id) + if meta: + existing = data.get("custom_metadata") or {} + existing.update(meta) + data["custom_metadata"] = existing + except Exception: + logger.warning( + "Failed to get KV transfer metadata for %s", + req_id, + exc_info=True, + ) + + mgr = self._ensure_kv_transfer_manager() + self._kv_extracted_req_ids = mgr.handle_finished_requests_kv_transfer( + finished_reqs=finished, + kv_caches=kv_caches, + block_size=self.cache_config.block_size, + cache_dtype=str(self.cache_config.cache_dtype), + ) + + +# ====================================================================== +# OmniAsyncOutput — async D2H for Omni AR outputs +# ====================================================================== + + +def _async_copy_to_np(x: torch.Tensor) -> np.ndarray: + return x.to("cpu", non_blocking=True).numpy() + + +def _async_copy_tensor(x: torch.Tensor) -> torch.Tensor: + return x.to("cpu", non_blocking=True) + + +def _async_copy_mm_value(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return _async_copy_tensor(value) + if isinstance(value, dict): + return {key: _async_copy_mm_value(val) for key, val in value.items()} + if isinstance(value, list): + return [_async_copy_mm_value(val) for val in value] + return value + + +def _async_copy_mm(mm_outputs: dict | None, total_tokens: int) -> dict[str, Any]: + """Non-blocking D2H copy of multimodal output tensors.""" + if not mm_outputs: + return {} + cpu: dict[str, Any] = {} + for k, v in mm_outputs.items(): + try: + cpu[k] = _async_copy_mm_value(v) + except Exception: + logger.exception("Error async-copying multimodal output %s", k) + return cpu + + +def _slice_pooler_value( + value: Any, + *, + req_index: int, + start: int, + end: int, + total_tokens: int, +) -> Any: + if isinstance(value, torch.Tensor): + if value.dim() > 0 and value.shape[0] == total_tokens: + return value[start:end].contiguous() + return value.clone() + if isinstance(value, dict): + return { + key: _slice_pooler_value( + val, + req_index=req_index, + start=start, + end=end, + total_tokens=total_tokens, + ) + for key, val in value.items() + } + if isinstance(value, list): + if not value: + return [] + elem = value[req_index] if req_index < len(value) else value[0] + return _slice_pooler_value( + elem, + req_index=req_index, + start=start, + end=end, + total_tokens=total_tokens, + ) + return value + + +def _ensure_tensor_values(payload: dict[str, Any]) -> dict[str, torch.Tensor]: + """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``. + + Non-tensor scalars (int/float/bool) are wrapped with ``torch.tensor()``; + values that cannot be safely converted are dropped. Enforces the tensor-only + invariant required by ``OmniEngineCoreOutput.multimodal_output`` (the channel + the async_chunk stage-input processor reads). Mirrors the V1 runner helper. + """ + result: dict[str, torch.Tensor] = {} + for key, val in payload.items(): + if isinstance(val, torch.Tensor): + result[key] = val + elif isinstance(val, (int, float, bool)): + result[key] = torch.tensor(val) + elif isinstance(val, (list, tuple)): + try: + result[key] = torch.tensor(val) + except (ValueError, TypeError, RuntimeError): + logger.warning( + "Dropping non-tensorizable multimodal output key '%s' (type=%s).", + key, + type(val).__name__, + ) + else: + logger.warning( + "Dropping non-tensor multimodal output key '%s' (type=%s).", + key, + type(val).__name__, + ) + return result + + +class OmniAsyncOutput(AsyncModelRunnerOutput): + """Async D2H copy for Omni AR model outputs. + + Mirrors upstream ``AsyncOutput`` but additionally handles + ``pooler_output`` (hidden states + multimodal outputs) via + non-blocking copies on the copy stream. + """ + + def __init__( + self, + model_runner_output: OmniModelRunnerOutput, + sampler_output: Any, + num_sampled_tokens: torch.Tensor, + main_stream: torch.cuda.Stream, + copy_stream: torch.cuda.Stream, + copy_event: torch.cuda.Event | None = None, + text_hidden: torch.Tensor | None = None, + multimodal_outputs: dict | None = None, + input_batch: Any | None = None, + ): + self.model_runner_output = model_runner_output + self.sampler_output = sampler_output + self.num_sampled_tokens = num_sampled_tokens + self.copy_event = copy_event if copy_event is not None else torch.cuda.Event() + + # Snapshot input_batch metadata needed for pooler_output slicing + self._need_pooler = text_hidden is not None + self._query_start_loc_np: np.ndarray | None = None + self._num_scheduled_tokens: np.ndarray | None = None + self._num_reqs: int = 0 + if self._need_pooler and input_batch is not None: + self._query_start_loc_np = input_batch.query_start_loc_np.copy() + self._num_scheduled_tokens = np.array(input_batch.num_scheduled_tokens, dtype=np.int32) + self._num_reqs = input_batch.num_reqs + + # Perform all D2H copies on the copy stream (non-blocking). + import contextlib + + @contextlib.contextmanager + def _stream(to_stream, from_stream): + try: + torch.cuda.set_stream(to_stream) + yield + finally: + torch.cuda.set_stream(from_stream) + + with _stream(copy_stream, main_stream): + copy_stream.wait_stream(main_stream) + + # Sampled token ids + self.sampled_token_ids_np = _async_copy_to_np(sampler_output.sampled_token_ids) + self.num_sampled_tokens_np = _async_copy_to_np(num_sampled_tokens) + + # Logprobs + self.logprobs_tensors = None + if sampler_output.logprobs_tensors is not None: + self.logprobs_tensors = sampler_output.logprobs_tensors.to_cpu_nonblocking() + self.num_nans: np.ndarray | None = None + if sampler_output.num_nans is not None: + self.num_nans = _async_copy_to_np(sampler_output.num_nans) + + # Prompt logprobs + self.prompt_logprobs_dict = { + k: v.to_cpu_nonblocking() if v is not None else None + for k, v in self.model_runner_output.prompt_logprobs_dict.items() + } + + # Pooler output (hidden + multimodal) — async D2H + self._hidden_cpu: torch.Tensor | None = None + self._mm_cpu: dict[str, Any] = {} + if self._need_pooler and text_hidden is not None: + self._hidden_cpu = _async_copy_tensor(text_hidden) + total_tokens = text_hidden.shape[0] + self._mm_cpu = _async_copy_mm(multimodal_outputs, total_tokens) + + self.copy_event.record(copy_stream) + + def get_output(self) -> OmniModelRunnerOutput: + self.copy_event.synchronize() + + # Sampled token ids + sampled_token_ids: list[list[int]] = self.sampled_token_ids_np.tolist() + num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist() + for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens): + del token_ids[num_tokens:] + self.model_runner_output.sampled_token_ids = sampled_token_ids + + # Logprobs + if self.num_nans is not None: + self.model_runner_output.num_nans_in_logits = dict( + zip(self.model_runner_output.req_ids, self.num_nans.tolist()) + ) + if self.logprobs_tensors is not None: + self.model_runner_output.logprobs = self.logprobs_tensors.tolists() + self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict + + # Pooler output. Populate two channels from the same per-request payloads, + # mirroring the V1 runner: + # * pooler_output -> sync/full-payload path (inline pooling_output bridge) + # * multimodal_outputs -> wire multimodal_output, which the async_chunk + # stage-input processor (talker2code2wav_async_chunk) reads for codes. + if self._need_pooler and self._hidden_cpu is not None: + pooler_output = OmniARModelRunner._build_pooler_output_from_cpu( + self._hidden_cpu, + self._mm_cpu, + self._query_start_loc_np, + self._num_scheduled_tokens, + self._num_reqs, + ) + self.model_runner_output.pooler_output = pooler_output + self.model_runner_output.multimodal_outputs = [ + _ensure_tensor_values(p) if p else {} for p in pooler_output + ] + + return self.model_runner_output diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py new file mode 100644 index 000000000..1c87b98cc --- /dev/null +++ b/vllm_omni/worker_v2/omni_generation_model_runner.py @@ -0,0 +1,398 @@ +"""OmniGenerationModelRunner — non-autoregressive stage runner on MR V2. + +Used for stages like Code2Wav that convert codec codes to audio waveforms. +No token sampling or logits computation — model output goes directly into +``pooler_output``. Inherits from ``OmniGPUModelRunner`` for intermediate +buffer and lifecycle hooks. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.worker.gpu.model_runner import ( + BatchDescriptor, + IntermediateTensors, + get_uniform_token_count, +) + +from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker_v2.forward_compat import add_forward_compat_kwargs +from vllm_omni.worker_v2.omni_ar_model_runner import _ensure_tensor_values +from vllm_omni.worker_v2.omni_model_runner import ( + OmniGPUModelRunner, + _make_execute_model_state, +) + +logger = init_logger(__name__) + + +class OmniGenerationModelRunner(OmniGPUModelRunner): + """Non-autoregressive generation runner (e.g. Code2Wav). + + Overrides ``execute_model`` to skip the tensor-only assertion and + ``sample_tokens`` to construct ``pooler_output`` from multimodal + model outputs without performing token sampling. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._gen_model_output: Any = None + self._gen_input_batch: Any = None + self._gen_kv_connector_output: Any = None + # Placeholder for ExecuteModelState.hidden_states — allocated + # once and reused every step to avoid per-forward allocation. + self._dummy_hidden = torch.zeros(1, dtype=self.dtype, device=self.device) + + # ------------------------------------------------------------------ + # Async-chunk support: replace prompt_token_ids for cached requests + # ------------------------------------------------------------------ + + def _handle_async_chunk_updates(self, scheduler_output: SchedulerOutput) -> None: + """In-place update cached requests whose prompt_token_ids changed. + + In async_chunk mode, the ``ChunkTransferAdapter`` replaces + ``Request.prompt_token_ids`` with new codec frames for each + chunk and resets ``num_computed_tokens`` to 0. The scheduler + propagates the new ``prompt_token_ids`` via + ``OmniCachedRequestData``. + + Instead of remove + re-add (which involves free_indices churn + and redundant model_state init), we update the existing slot + in-place. This is safe for Code2Wav because: + - No KV cache / rope state to reinitialize + - staged writes are applied once at the end + + The old intermediate buffer for this slot is cleared here; the + inherited ``OmniGPUModelRunner.update_requests`` (called right after + this method in ``execute_model``) writes the current chunk state. + """ + cached = scheduler_output.scheduled_cached_reqs + if not cached.req_ids: + return + + if not isinstance(cached, OmniCachedRequestData): + return + + new_prompt_ids = cached.prompt_token_ids + if not new_prompt_ids: + return + + updated = False + released_chunks: list[OmniNewRequestData] = [] + + for i, req_id in enumerate(cached.req_ids): + new_ids = new_prompt_ids.get(req_id) + if new_ids is None: + continue + + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + block_ids = cached.new_block_ids[i] + released_chunks.append( + OmniNewRequestData( + req_id=req_id, + prompt_token_ids=new_ids, + mm_features=[], + sampling_params=None, + pooling_params=None, + block_ids=block_ids if block_ids is not None else tuple(), + num_computed_tokens=0, + lora_request=None, + prompt_embeds=None, + prefill_token_ids=new_ids, + additional_information=cached.additional_information.get(req_id), + ) + ) + continue + + self.model_state.intermediate_buffer.remove_request(req_idx) + + # In-place update token state — same slot, no remove/re-add. + # .np[] = direct write (no GPU buffer); stage_write = GPU-synced. + n = len(new_ids) + self.req_states.prompt_len.np[req_idx] = n + self.req_states.prefill_len.np[req_idx] = n + self.req_states.total_len.stage_write_elem(req_idx, n) + self.req_states.all_token_ids.stage_write(req_idx, 0, new_ids) + self.req_states.num_computed_tokens.stage_write_elem(req_idx, 0) + self.req_states.num_computed_prefill_tokens[req_idx] = 0 + + updated = True + + if released_chunks: + self.add_requests(SimpleNamespace(scheduled_new_reqs=released_chunks)) + if updated: + self.req_states.apply_staged_writes() + + def _release_generation_slots(self, input_batch: Any) -> None: + if not getattr(self.model_config, "async_chunk", False): + return + model_state = getattr(self, "model_state", None) + remove_request = getattr(self, "_remove_request", None) + if model_state is None or remove_request is None: + return + for i in range(input_batch.num_reqs): + req_id = input_batch.req_ids[i] + req_idx = int(input_batch.idx_mapping_np[i]) + model_state.remove_request(req_idx) + remove_request(req_id) + + # ------------------------------------------------------------------ + # profile / warmup — skip sampler since there are no logits + # ------------------------------------------------------------------ + + @torch.inference_mode() + def profile_run(self) -> None: + """Generation models have no KV cache — skip profiling. + + Code2Wav shares GPU memory with the Talker stage (same device); + its memory footprint is managed via ``gpu_memory_utilization`` + config, not profiled dynamically. Running the real model with + random input_ids causes out-of-bounds indexing in codec lookup + tables. + """ + torch.accelerator.synchronize() + + # ------------------------------------------------------------------ + # execute_model — run the generation model, store raw output + # ------------------------------------------------------------------ + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + ) -> ModelRunnerOutput | IntermediateTensors | None: + if not dummy_run: + self.finish_requests(scheduler_output) + self.free_states(scheduler_output) + # Handle async_chunk prompt_token_ids replacement for cached + # requests BEFORE add/update — update the existing slot + # in-place with the new chunk's tokens. + self._handle_async_chunk_updates(scheduler_output) + self.add_requests(scheduler_output) + self.update_requests(scheduler_output) + self.block_tables.apply_staged_writes() + if scheduler_output.total_num_scheduled_tokens == 0: + return self.kv_connector.no_forward(scheduler_output) + + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len) + batch_desc, _ = self._dispatch_batch_descriptor( + num_reqs=num_reqs, + num_toks=num_toks, + uniform_tok_count=uniform_tok_count, + use_eager=is_profile, + ) + + if batch_desc.num_tokens == 0: + return self.kv_connector.no_forward(scheduler_output) + + if not dummy_run: + input_batch = self.prepare_inputs(scheduler_output, batch_desc) + else: + from vllm.v1.worker.gpu.input_batch import InputBatch + + input_batch = InputBatch.make_dummy( + batch_desc.num_reqs or num_reqs, + batch_desc.num_tokens, + self.input_buffers, + ) + + attn_metadata = None + slot_mappings_by_layer = None + + inputs_embeds = None + if self.supports_mm_inputs and self.is_first_pp_rank: + # vLLM 0.23 dropped the req_states arg from get_mm_embeddings. + inputs_embeds = self.model_state.get_mm_embeddings( + scheduler_output.scheduled_encoder_inputs, + input_batch, + ) + + model_inputs = { + "input_ids": input_batch.input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": intermediate_tensors, + **self.model_state.prepare_inputs(input_batch, self.req_states), + } + add_forward_compat_kwargs(model_inputs, input_batch, self.sampler) + + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=self.lora_config is not None, + ) + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=batch_desc.cg_mode, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + ): + self.kv_connector.pre_forward(scheduler_output) + model_output = self.model(**model_inputs) + + kv_connector_output = self.kv_connector.post_forward(scheduler_output.finished_req_ids) + + # Convert raw model output to OmniOutput. + if not isinstance(model_output, OmniOutput): + buffer_list = self.model_state.intermediate_buffer.gather(input_batch) + try: + model_output = self.model.make_omni_output( + model_output, + model_intermediate_buffer=buffer_list, + runtime_additional_information=buffer_list, + ) + except Exception: + logger.error( + "make_omni_output failed; returning empty output", + exc_info=True, + ) + self._gen_model_output = None + self.execute_model_state = _make_execute_model_state( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=self._dummy_hidden, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + kv_connector_output=kv_connector_output, + num_tokens_across_dp=None, + ) + return None + + self._gen_model_output = model_output + self._gen_input_batch = input_batch + self._gen_kv_connector_output = kv_connector_output + + # ExecuteModelState is required by the upstream engine loop + # (EngineCore checks execute_model_state is not None before + # calling sample_tokens). + self.execute_model_state = _make_execute_model_state( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=self._dummy_hidden, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + kv_connector_output=kv_connector_output, + num_tokens_across_dp=None, + ) + return None + + # ------------------------------------------------------------------ + # sample_tokens — build pooler_output, no actual sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() + def sample_tokens(self, grammar_output: GrammarOutput | None = None) -> OmniModelRunnerOutput | None: + model_output = self._gen_model_output + input_batch = self._gen_input_batch + kv_connector_output = self._gen_kv_connector_output + self._gen_model_output = None + self._gen_input_batch = None + self._gen_kv_connector_output = None + self.execute_model_state = None + + if model_output is None or input_batch is None: + return None + + num_reqs = input_batch.num_reqs + + # Mark all scheduled tokens as computed so the scheduler does + # not re-schedule them. Unlike AR stages we do NOT call + # self.postprocess() — that kernel advances num_computed_tokens + # by 1 and emits sampled tokens, which would cause check_stop + # to fire. Instead, set num_computed_tokens = prompt_len + # directly, matching V1's behavior. + for i in range(num_reqs): + req_idx = int(input_batch.idx_mapping_np[i]) + prompt_len = int(self.req_states.prompt_len.np[req_idx]) + self.req_states.num_computed_tokens.stage_write_elem(req_idx, prompt_len) + self.req_states.num_computed_tokens.apply_write() + + # Build pooler_output from OmniOutput.multimodal_outputs (dict). + pooler_output = self._build_pooler_output(model_output, num_reqs) + + req_ids = input_batch.req_ids + + # Generation models don't do token sampling. Return one empty + # list per request so the scheduler does NOT trigger check_stop + # (which would prematurely finish the request). The request + # stays RUNNING until the orchestrator marks it done via + # chunk_transfer_adapter.finished_requests. + sampled_token_ids: list[list[int]] = [[] for _ in range(len(req_ids))] + self._release_generation_slots(input_batch) + + # Per-request multimodal_output (tensor-only) built from the same payloads. + # OmniGenerationScheduler indexes this as mm_outputs[req_index], and for the + # final (code2wav) stage it becomes the user-facing outputs[0].multimodal_output. + # Must be a per-request list (not the raw dict), matching the V1 runner. + multimodal_outputs = [_ensure_tensor_values(p) if p else {} for p in pooler_output] + + return OmniModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={rid: i for i, rid in enumerate(req_ids)}, + sampled_token_ids=sampled_token_ids, + pooler_output=pooler_output, + multimodal_outputs=multimodal_outputs, + kv_connector_output=kv_connector_output, + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _build_pooler_output( + model_output: OmniOutput, + num_reqs: int, + ) -> list[dict[str, Any] | None]: + """Extract per-request pooler payloads from model output. + + Code2Wav's ``make_omni_output`` returns + ``{"model_outputs": [tensor_per_req, ...], "sr": [...]}``, + so each value is a ``list`` with ``len == num_reqs``. + """ + mm = model_output.multimodal_outputs + if not isinstance(mm, dict): + logger.warning( + "Unexpected multimodal_outputs type: %s; returning empty pooler_output", + type(mm).__name__ if mm is not None else "None", + ) + return [None] * num_reqs + + pooler: list[dict[str, Any] | None] = [] + for i in range(num_reqs): + payload: dict[str, Any] = {} + for key, val in mm.items(): + # Primary path: val is list[Tensor] with len == num_reqs + # (Code2Wav make_omni_output format). + if isinstance(val, list) and len(val) == num_reqs: + out = val[i] + payload[key] = out.detach().cpu().contiguous() if isinstance(out, torch.Tensor) else out + elif isinstance(val, torch.Tensor): + if val.dim() > 0 and val.shape[0] == num_reqs: + payload[key] = val[i].detach().cpu().contiguous() + else: + payload[key] = val.detach().cpu().contiguous() + else: + payload[key] = val + pooler.append(payload) + return pooler diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py new file mode 100644 index 000000000..2c990dd8e --- /dev/null +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -0,0 +1,393 @@ +"""Omni v2 GPU model runner hooks.""" + +from __future__ import annotations + +import threading +from typing import Any + +import torch +from vllm.config.compilation import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.worker.gpu.model_runner import ( + BatchDescriptor, + BatchExecutionDescriptor, + ExecuteModelState, + GPUModelRunner, + IntermediateTensors, + build_slot_mappings_by_layer, + get_uniform_token_count, +) + +from vllm_omni.compat import make_filtered_namedtuple +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.forward_compat import add_forward_compat_kwargs +from vllm_omni.worker_v2.model_states import init_omni_model_state +from vllm_omni.worker_v2.model_states.intermediate_buffer import ( + _resolve_additional_information, +) +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +logger = init_logger(__name__) + +_model_state_patch_lock = threading.RLock() +_GPU_MODEL_RUNNER_HAS_SHUTDOWN = hasattr(GPUModelRunner, "shutdown") +# Fields the omni runners still pass to ExecuteModelState that current vLLM no +# longer declares on the namedtuple. They are filtered out by +# make_filtered_namedtuple and carried out-of-band instead: +# - num_tokens_across_dp: only consumed inside execute_model's forward context. +# - kv_connector_output: removed from ExecuteModelState in vLLM 0.23.0 (post_forward +# now runs in sample_tokens). Carried via self._last_kv_connector_output (AR) / +# self._gen_kv_connector_output (generation). +_KNOWN_EXECUTE_MODEL_STATE_COMPAT_FIELDS = {"num_tokens_across_dp", "kv_connector_output"} + + +def _make_execute_model_state(**kwargs): + state, unknown = make_filtered_namedtuple( + ExecuteModelState, + known_extra_fields=_KNOWN_EXECUTE_MODEL_STATE_COMPAT_FIELDS, + **kwargs, + ) + if unknown: + logger.warning("Unknown fields passed to ExecuteModelState: %s", sorted(unknown)) + return state + + +def _needs_capture_tensor_unwrap(model: Any) -> bool: + return bool(getattr(model, "_returns_tuple", False) or getattr(model, "model_stage", None) == "thinker") + + +class OmniGPUModelRunner(GPUModelRunner): + """Thin layer over v2 ``GPUModelRunner`` for Omni lifecycle hooks.""" + + model_state: OmniModelState + _last_aux_output: Any + _last_multimodal_outputs: dict[str, Any] | None + _model_returns_tuple: bool + + def shutdown(self) -> None: + if _GPU_MODEL_RUNNER_HAS_SHUTDOWN: + super().shutdown() + + def load_model(self, *args: Any, **kwargs: Any) -> None: + import vllm.v1.worker.gpu.model_runner as _mr_module + + with _model_state_patch_lock: + _orig = _mr_module.init_model_state + _mr_module.init_model_state = init_omni_model_state + try: + super().load_model(*args, **kwargs) + finally: + _mr_module.init_model_state = _orig + self._last_aux_output = None + self._last_multimodal_outputs = None + self._last_kv_connector_output = None + self._model_returns_tuple = _needs_capture_tensor_unwrap(self.model) + self._exclude_full_graph = self._model_returns_tuple or hasattr(self.model, "_last_captured_layers") + + # Preprocess models own embedding buffers; encoder_runner sizing would mismatch. + if getattr(self.model, "has_preprocess", False) and self.supports_mm_inputs: + self.supports_mm_inputs = False + self.encoder_cache = None + + # ------------------------------------------------------------------ + # CUDA Graph: conditionally exclude FULL mode + # ------------------------------------------------------------------ + + def capture_model(self) -> int: + """Handle CUDA graph capture for Omni models. + + Exclude FULL mode for tuple-returning models because + ``run_fullgraph`` bypasses Python-level tuple intercept. + + For PIECEWISE capture, the warmup pass runs with + ``CUDAGraphMode.NONE`` which hits ``torch.empty_like(hidden_states)`` + in the cudagraph framework. If the model returns a tuple, that call + crashes. We temporarily wrap the model's forward to extract only the + tensor part during capture, then restore the original forward. + """ + if self._exclude_full_graph: + mgr = self.cudagraph_manager + # NOTE: ``_capture_descs`` / ``_candidates`` are private CudaGraphManager + # internals (present and same-typed as of vLLM 0.23.0). Guard with hasattr + # so a future vLLM that renames/removes them degrades to a warning instead + # of an AttributeError; the worst case is FULL graphs not being excluded. + if hasattr(mgr, "_capture_descs") and hasattr(mgr, "_candidates"): + if CUDAGraphMode.FULL in mgr._capture_descs: + del mgr._capture_descs[CUDAGraphMode.FULL] + for i, descs in enumerate(mgr._candidates): + mgr._candidates[i] = [d for d in descs if d.cg_mode != CUDAGraphMode.FULL] + logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.") + else: + logger.warning( + "CudaGraphManager lacks _capture_descs/_candidates; cannot exclude " + "FULL graph capture for tuple-returning Omni model. vLLM internals " + "may have changed — verify omni_model_runner.capture_model()." + ) + + # Wrap model forward during capture so tuple returns don't crash + # torch.empty_like() in the PIECEWISE warmup pass. + if self._model_returns_tuple: + original_forward = self.model.forward + + def _capture_forward(*args: Any, **kwargs: Any) -> torch.Tensor: + output = original_forward(*args, **kwargs) + if isinstance(output, OmniOutput): + return output.text_hidden_states + if isinstance(output, tuple): + return output[0] + return output + + self.model.forward = _capture_forward # type: ignore[assignment] + try: + return super().capture_model() + finally: + self.model.forward = original_forward # type: ignore[assignment] + return super().capture_model() + + def _dispatch_batch_descriptor( + self, + *, + num_reqs: int, + num_toks: int, + uniform_tok_count: int, + use_eager: bool, + ): + if use_eager: + batch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_toks, + num_reqs=num_reqs, + ) + else: + batch_desc = self.cudagraph_manager.dispatch(num_reqs, num_toks, uniform_tok_count) + if self.dp_size > 1: + from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding + + return sync_cudagraph_and_dp_padding( + self.cudagraph_manager, + batch_desc, + num_toks, + num_reqs, + uniform_tok_count, + self.dp_size, + self.dp_rank, + ) + return batch_desc, None + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, + ) -> Any: + if not dummy_run: + self.finish_requests(scheduler_output) + self.free_states(scheduler_output) + self.add_requests(scheduler_output) + self.update_requests(scheduler_output) + self.block_tables.apply_staged_writes() + if scheduler_output.total_num_scheduled_tokens == 0: + return self.kv_connector.no_forward(scheduler_output) + + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_toks = scheduler_output.total_num_scheduled_tokens + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len) + # Encoder-decoder models: disable compilation when encoder inputs + # are scheduled (dynamic cross-attention cache updates). + skip_compiled = self.is_encoder_decoder and bool(scheduler_output.scheduled_encoder_inputs) + batch_desc, num_tokens_across_dp = self._dispatch_batch_descriptor( + num_reqs=num_reqs, + num_toks=num_toks, + uniform_tok_count=uniform_tok_count, + use_eager=is_profile or skip_compiled, + ) + + if batch_desc.num_tokens == 0: + return self.kv_connector.no_forward(scheduler_output) + + if not dummy_run: + input_batch = self.prepare_inputs(scheduler_output, batch_desc) + block_tables, slot_mappings = self.prepare_attn(input_batch) + + if self.lora_config: + lora_inputs = self.lora_state.make_lora_inputs( + input_batch.req_ids, + input_batch.idx_mapping_np, + input_batch.num_scheduled_tokens, + ) + self._set_active_loras(*lora_inputs) + else: + from vllm.v1.worker.gpu.input_batch import InputBatch + + input_batch = InputBatch.make_dummy( + batch_desc.num_reqs or num_reqs, + batch_desc.num_tokens, + self.input_buffers, + ) + if not skip_attn_for_dummy_run: + block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) + else: + block_tables = None + slot_mappings = None + + attn_metadata = None + slot_mappings_by_layer = None + if not (dummy_run and skip_attn_for_dummy_run): + assert slot_mappings is not None + slot_mappings_by_layer = build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config) + assert block_tables is not None + attn_metadata = self.model_state.prepare_attn( + input_batch, + batch_desc.cg_mode, + block_tables, + slot_mappings, + self.attn_groups, + self.kv_cache_config, + ) + + inputs_embeds = None + if self.supports_mm_inputs and self.is_first_pp_rank: + # vLLM 0.23 dropped the req_states arg from get_mm_embeddings. + inputs_embeds = self.model_state.get_mm_embeddings( + scheduler_output.scheduled_encoder_inputs, + input_batch, + ) + + model_inputs: dict[str, Any] = { + "input_ids": input_batch.input_ids, + "positions": input_batch.positions, + "inputs_embeds": inputs_embeds, + "intermediate_tensors": intermediate_tensors, + **self.model_state.prepare_inputs(input_batch, self.req_states), + } + add_forward_compat_kwargs(model_inputs, input_batch, self.sampler) + if not self.is_first_pp_rank: + model_inputs["input_ids"] = None + model_inputs["inputs_embeds"] = None + assert intermediate_tensors is not None + + # ★ PRE-FORWARD: per-request preprocess + batched MTP. + # Runs for ALL graph modes (FULL, PIECEWISE, NONE). + # For FULL graph: OmniModelState provides a static inputs_embeds + # buffer that was captured by the graph. Preprocess writes + # in-place to this buffer, and FULL graph replay reads the + # updated values from the same tensor address. + if not dummy_run: + self.model_state.run_preprocess(input_batch, model_inputs) + + # --- Model forward --- + if batch_desc.cg_mode == CUDAGraphMode.FULL: + # FULL graph replay. Preprocess already wrote to the static + # inputs_embeds buffer above. + assert self.cudagraph_manager is not None + self.kv_connector.pre_forward(scheduler_output) + model_output = self.cudagraph_manager.run_fullgraph(batch_desc) + hidden_states = model_output + self._last_aux_output = None + self._last_multimodal_outputs = None + else: + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=self.lora_config is not None, + ) + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=batch_desc.cg_mode, + num_tokens_across_dp=num_tokens_across_dp, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + skip_compiled=skip_compiled, + ): + self.kv_connector.pre_forward(scheduler_output) + model_output = self.model(**model_inputs) + + # Extract hidden_states from model output. + self._last_aux_output = None + self._last_multimodal_outputs = None + if isinstance(model_output, OmniOutput): + hidden_states = model_output.text_hidden_states + if model_output.multimodal_outputs: + self._last_multimodal_outputs = model_output.multimodal_outputs + elif isinstance(model_output, tuple) and len(model_output) == 2: + hidden_states, self._last_aux_output = model_output + if hasattr(self.model, "_last_captured_layers"): + self.model._last_captured_layers = self._last_aux_output + else: + hidden_states = model_output + + if not dummy_run and isinstance(hidden_states, torch.Tensor): + self.model_state.run_postprocess(hidden_states, input_batch) + + kv_connector_output = self.kv_connector.post_forward(scheduler_output.finished_req_ids) + self._last_kv_connector_output = kv_connector_output + self.execute_model_state = _make_execute_model_state( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + hidden_states=hidden_states, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + kv_connector_output=kv_connector_output, + num_tokens_across_dp=num_tokens_across_dp, + ) + + if not self.is_last_pp_rank: + assert isinstance(hidden_states, IntermediateTensors) + hidden_states.kv_connector_output = kv_connector_output + return hidden_states + assert isinstance(hidden_states, torch.Tensor) + return None + + # ------------------------------------------------------------------ + # Request lifecycle: update intermediate buffer from cached requests + # ------------------------------------------------------------------ + + def update_requests(self, scheduler_output: SchedulerOutput) -> None: + """Merge updated additional_information into intermediate_buffer. + + In async_chunk mode, chunk_transfer_adapter attaches updated + additional_information (e.g. thinker_decode_embeddings) to + OmniCachedRequestData for cached requests every schedule step. + Upstream GPUModelRunner.update_requests does not handle this + field, so we merge it into the intermediate buffer here. + """ + super().update_requests(scheduler_output) + + cached = scheduler_output.scheduled_cached_reqs + addl_info = getattr(cached, "additional_information", None) + if not addl_info: + return + for req_id, info in addl_info.items(): + if info is None: + continue + req_idx = self.req_states.req_id_to_index.get(req_id) + if req_idx is None: + continue + resolved = _resolve_additional_information(info) + if resolved: + self.model_state.intermediate_buffer.update(req_idx, resolved) + + # ------------------------------------------------------------------ + # Request lifecycle: clean up intermediate buffer on finish + # ------------------------------------------------------------------ + + def finish_requests(self, scheduler_output: SchedulerOutput) -> None: + # IMPORTANT: Must query req_id_to_index BEFORE super().finish_requests() + # because super() calls req_states.remove_request(req_id) which pops the + # mapping and returns the slot index to free_indices. + finished = scheduler_output.finished_req_ids + preempted = scheduler_output.preempted_req_ids + all_done = finished | preempted if preempted else finished + for req_id in all_done: + idx = self.req_states.req_id_to_index.get(req_id) + if idx is not None: + self.model_state.remove_request(idx) + super().finish_requests(scheduler_output) diff --git a/vllm_omni/worker_v2/outputs.py b/vllm_omni/worker_v2/outputs.py new file mode 100644 index 000000000..89fcc6ad8 --- /dev/null +++ b/vllm_omni/worker_v2/outputs.py @@ -0,0 +1,5 @@ +"""Re-export OmniModelRunnerOutput for worker_v2 consumers.""" + +from vllm_omni.outputs import OmniModelRunnerOutput + +__all__ = ["OmniModelRunnerOutput"]