diff --git a/benchmarks/qwen3_omni/embedding_harness.py b/benchmarks/qwen3_omni/embedding_harness.py new file mode 100644 index 000000000..63d14ef0d --- /dev/null +++ b/benchmarks/qwen3_omni/embedding_harness.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Small fixtures for the real Qwen3 thinker embedding method and its oracle.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from types import SimpleNamespace + +import torch +from torch import nn +from torch.nn import functional as F + +SCENARIOS = ("text", "empty", "audio", "image", "video", "mixed", "interleaved", "vision_no_deepstack") + + +@dataclass(frozen=True) +class ModelShape: + vocab_size: int = 152064 + hidden_size: int = 2048 + deepstack_levels: int = 3 + image_token_id: int = 151655 + video_token_id: int = 151656 + audio_token_id: int = 151675 + + def __post_init__(self): + if min(self.vocab_size, self.hidden_size, self.deepstack_levels) <= 0: + raise ValueError("vocab_size, hidden_size and deepstack_levels must be positive") + ids = (self.image_token_id, self.video_token_id, self.audio_token_id) + if len(set(ids)) != 3 or not all(0 < token < self.vocab_size for token in ids): + raise ValueError("Provide three distinct, positive, in-vocabulary modality IDs; use oov=True to test OOV") + + def token_ids(self, oov=False): + if oov: + return dict(zip(("image", "video", "audio"), range(self.vocab_size, self.vocab_size + 3))) + return {name: getattr(self, f"{name}_token_id") for name in ("image", "video", "audio")} + + +@dataclass +class Feature: + modality: str + positions: torch.Tensor + main: torch.Tensor + deepstack: torch.Tensor | None + combined: torch.Tensor + + +@dataclass +class EmbeddingCase: + name: str + input_ids: torch.Tensor + is_multimodal: torch.Tensor | None + features: tuple[Feature, ...] + fingerprint: str + + def fresh_embeddings(self): + # The thinker replaces list entries with main-scale views. + return None if self.name == "text" else [feature.combined for feature in self.features] + + +def _tensor_bytes(tensor): + return tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + + +def build_case( + name, + num_tokens, + num_mm_tokens, + shape, + *, + device="cuda", + dtype=torch.bfloat16, + mask_device="cpu", + seed=42, + oov=False, +): + """Build known feature placements; no model merge is used to define them.""" + if name not in SCENARIOS: + raise ValueError(f"Unknown scenario: {name}") + if num_tokens <= 0 or not 0 <= num_mm_tokens < num_tokens: + raise ValueError("Require N > 0 and 0 <= M < N") + if (name in {"text", "empty"}) != (num_mm_tokens == 0): + raise ValueError("Only text/empty scenarios accept M=0") + if name in {"mixed", "interleaved"} and (num_mm_tokens < 4 or num_mm_tokens % 2): + raise ValueError("Mixed/interleaved scenarios require an even M >= 4") + token_ids = shape.token_ids(oov) + generator = torch.Generator().manual_seed(seed) + ids = torch.randint(0, min(shape.vocab_size, *token_ids.values()), (num_tokens,), generator=generator) + mask = torch.zeros(num_tokens, dtype=torch.bool) + positions = torch.arange((num_tokens - num_mm_tokens) // 2, (num_tokens - num_mm_tokens) // 2 + num_mm_tokens) + groups = [] + if name == "mixed": + groups = [("audio", positions[: num_mm_tokens // 2]), ("image", positions[num_mm_tokens // 2 :])] + elif name == "interleaved": + groups = [("video", positions[::2]), ("audio", positions[1::2])] + elif num_mm_tokens: + groups = [("image" if name == "vision_no_deepstack" else name, positions)] + features = [] + for modality, indexes in groups: + ids[indexes] = token_ids[modality] + mask[indexes] = True + rows = torch.arange(len(indexes), dtype=torch.float32)[:, None].remainder(31) / 32 + cols = torch.arange(shape.hidden_size, dtype=torch.float32)[None, :].remainder(7) / 16 + main_cpu = rows + cols + {"audio": 2, "image": 4, "video": 6}[modality] + main = main_cpu.to(device=device, dtype=dtype) + deepstack = None + parts = [main] + if modality != "audio" and name != "vision_no_deepstack": + deepstack = torch.stack( + [(main_cpu + 16 + 4 * level).to(device=device, dtype=dtype) for level in range(shape.deepstack_levels)] + ) + parts.extend(deepstack.unbind(0)) + combined = torch.cat(parts, dim=-1) + # set_mm_embedding_modality simply attaches this attribute. Tag only + # after cat/to: views and copies do not preserve arbitrary attributes. + combined.modality = modality + features.append(Feature(modality, indexes, main, deepstack, combined)) + digest = hashlib.sha256(json.dumps({"name": name, "shape": asdict(shape), "seed": seed, "oov": oov}).encode()) + digest.update(_tensor_bytes(ids)) + digest.update(_tensor_bytes(mask)) + for feature in features: + digest.update(feature.modality.encode()) + digest.update(_tensor_bytes(feature.combined)) + return EmbeddingCase( + name, + ids.to(device=device), + None if name == "text" else mask.to(device=mask_device), + tuple(features), + digest.hexdigest(), + ) + + +class LanguageModelAdapter(nn.Module): + def __init__(self, embedding): + super().__init__() + self.embedding = embedding + + def embed_input_ids(self, input_ids): + return self.embedding(input_ids) + + def forward(self, input_ids): + return self.embed_input_ids(input_ids) + + +def make_thinker( + shape, + *, + device="cuda", + dtype=torch.bfloat16, + seed=42, + backend="torch", + deepstack=True, + buffer_capacity=32, + oov=False, +): + """Skip checkpoint construction, preserving the real class and its MRO.""" + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + ) + + if buffer_capacity <= 0: + raise ValueError("buffer_capacity must be positive") + device = torch.device(device) + if backend == "torch": + embedding = nn.Embedding(shape.vocab_size, shape.hidden_size, device=device, dtype=dtype) + elif backend == "vllm": + from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + with torch.device(device): + embedding = VocabParallelEmbedding(shape.vocab_size, shape.hidden_size, params_dtype=dtype, disable_tp=True) + else: + raise ValueError(f"Unknown embedding backend: {backend}") + with torch.no_grad(): + embedding.weight.uniform_(-0.5, 0.5, generator=torch.Generator(device=device).manual_seed(seed)) + thinker = object.__new__(Qwen3OmniMoeThinkerForConditionalGeneration) + nn.Module.__init__(thinker) + thinker.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=shape.hidden_size)) + for modality, token_id in shape.token_ids(oov).items(): + setattr(thinker.config, f"{modality}_token_id", token_id) + thinker.visual = SimpleNamespace( + deepstack_visual_indexes=list(range(shape.deepstack_levels)) if deepstack else None + ) + thinker.language_model = LanguageModelAdapter(embedding) + thinker._has_oov_mm_tokens = oov + thinker.use_deepstack = deepstack + thinker.deepstack_num_level = shape.deepstack_levels if deepstack else 0 + thinker.deepstack_input_embeds = [ + torch.zeros(buffer_capacity, shape.hidden_size, device=device, dtype=dtype) + for _ in range(thinker.deepstack_num_level) + ] + thinker.deepstack_input_embeds_num_tokens = 0 + return thinker.eval() + + +@torch.inference_mode() +def expected_outputs(model, case): + """Independent oracle: direct lookup and explicit positions, not merge helpers.""" + ids = case.input_ids + if model._has_oov_mm_tokens and case.is_multimodal is not None: + ids = ids.masked_fill(case.is_multimodal.to(ids.device), 0) + main = F.embedding(ids, model.language_model.embedding.weight) + has_deepstack = any(feature.deepstack is not None for feature in case.features) + deepstack = main.new_zeros((model.deepstack_num_level, *main.shape)) if has_deepstack else None + for feature in case.features: + indexes = feature.positions.to(main.device) + main[indexes] = feature.main + if feature.deepstack is not None: + deepstack[:, indexes, :] = feature.deepstack + return main, deepstack + + +@torch.inference_mode() +def assert_case_output(model, case, result): + expected, deepstack = expected_outputs(model, case) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + if deepstack is None: + assert model.deepstack_input_embeds_num_tokens == 0, "Unexpected stale DeepStack buffer" + else: + assert model.deepstack_input_embeds_num_tokens == case.input_ids.numel() + actual = torch.stack([buffer[: case.input_ids.numel()] for buffer in model.deepstack_input_embeds]) + torch.testing.assert_close(actual, deepstack, rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py new file mode 100644 index 000000000..d041eace0 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_embed.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Run Qwen3-Omni embedding regressions on a Linux/CUDA environment. + +The oracle assigns features to token positions independently of the model's +merge helpers. LM call counts additionally catch issue #7451: ordinary +multimodal inputs must not repeat the text embedding lookup. Small synthetic +weights suffice; no checkpoint is downloaded or loaded. +""" + +import sys + +import pytest +import torch +from pytest_mock import MockerFixture + +from benchmarks.qwen3_omni.embedding_harness import ( + ModelShape, + assert_case_output, + build_case, + expected_outputs, + make_thinker, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +_SHAPE = ModelShape( + vocab_size=64, + hidden_size=8, + deepstack_levels=3, + image_token_id=60, + video_token_id=61, + audio_token_id=62, +) + +_SCENARIOS = ( + "text", + "empty", + "audio", + "image", + "video", + "mixed", + "interleaved", + "vision_no_deepstack", +) + + +@pytest.fixture(scope="module", autouse=True) +def require_cuda_environment(): + if sys.platform != "linux" or torch.version.cuda is None or not torch.cuda.is_available(): + pytest.fail( + "This validation suite requires a Linux/CUDA environment; " + "CPU execution and an all-skipped result are not valid substitutes.", + pytrace=False, + ) + if not torch.cuda.is_bf16_supported(including_emulation=False): + pytest.fail("This validation suite requires native BF16 support on the current CUDA device.", pytrace=False) + + +@torch.inference_mode() +def _embed(model, case): + return model.embed_input_ids( + case.input_ids, + multimodal_embeddings=case.fresh_embeddings(), + is_multimodal=case.is_multimodal, + ) + + +@pytest.mark.parametrize( + ("scenario", "dtype"), + [pytest.param(scenario, torch.bfloat16, id=f"{scenario}-bf16") for scenario in _SCENARIOS] + + [pytest.param(scenario, torch.float32, id=f"{scenario}-fp32") for scenario in ("image", "mixed", "interleaved")], +) +def test_embedding_positions_and_single_lm_lookup(mocker: MockerFixture, scenario: str, dtype: torch.dtype): + model = make_thinker(_SHAPE, device="cuda", dtype=dtype, deepstack=scenario != "vision_no_deepstack") + case = build_case( + scenario, + 16, + 0 if scenario in {"text", "empty"} else 8, + _SHAPE, + device="cuda", + dtype=dtype, + mask_device="cpu", + ) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + # Unlike upstream-only mocks, make_thinker returns an instance of Omni's + # actual thinker class, so its zero-argument super() follows the real MRO. + lm_embed.assert_called_once() + assert_case_output(model, case, result) + + +@pytest.mark.parametrize("scenario", ["image", "mixed", "interleaved"]) +def test_repeated_calls_keep_original_deepstack_features(mocker: MockerFixture, scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + original_embeddings = case.fresh_embeddings() + assert original_embeddings + original_shapes = [tensor.shape for tensor in original_embeddings] + original_modalities = [getattr(tensor, "modality", None) for tensor in original_embeddings] + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + previous_list = None + + for _ in range(3): + embeddings = case.fresh_embeddings() + assert embeddings is not original_embeddings + assert embeddings is not previous_list + assert [tensor.shape for tensor in embeddings] == original_shapes + assert [getattr(tensor, "modality", None) for tensor in embeddings] == original_modalities + lm_embed.reset_mock() + + with torch.inference_mode(): + result = model.embed_input_ids( + case.input_ids, + multimodal_embeddings=embeddings, + is_multimodal=case.is_multimodal, + ) + + lm_embed.assert_called_once() + assert_case_output(model, case, result) + # torch.split replaces entries in the supplied list with main-scale + # views. The next call must receive the original multiscale tensors. + assert [tensor.shape for tensor in original_embeddings] == original_shapes + model._clear_deepstack_input_embeds(case.input_ids.numel()) + previous_list = embeddings + + +@pytest.mark.parametrize("scenario", ["audio", "image", "mixed", "interleaved"]) +def test_oov_multimodal_ids_are_masked_before_lm_lookup(mocker: MockerFixture, scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, oov=True) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu", oov=True) + original_ids = case.input_ids.clone() + mask = case.is_multimodal.to(device=case.input_ids.device) + assert torch.all(case.input_ids[mask] >= _SHAPE.vocab_size) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + lm_embed.assert_called_once() + received_ids = lm_embed.call_args.args[0] + torch.testing.assert_close(received_ids, original_ids.masked_fill(mask, 0), rtol=0, atol=0) + torch.testing.assert_close(case.input_ids, original_ids, rtol=0, atol=0) + assert_case_output(model, case, result) + + +@pytest.mark.parametrize("next_scenario", ["text", "empty"]) +def test_padded_deepstack_read_and_clear_do_not_leak_into_decode(mocker: MockerFixture, next_scenario: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, buffer_capacity=32) + image_case = build_case("image", 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + image_result = _embed(model, image_case) + assert_case_output(model, image_case, image_result) + _, expected_deepstack = expected_outputs(model, image_case) + assert expected_deepstack is not None + + num_valid = image_case.input_ids.numel() + num_padded = 24 + # Padding may contain old buffer contents. _get must zero the tail while + # retaining all valid visual features, as the real forward requires. + for buffer in model.deepstack_input_embeds: + buffer[num_valid:num_padded].fill_(77) + padded = model._get_deepstack_input_embeds(num_padded) + assert padded is not None + for level in range(_SHAPE.deepstack_levels): + actual = padded[f"deepstack_input_embeds_{level}"] + torch.testing.assert_close(actual[:num_valid], expected_deepstack[level]) + assert torch.count_nonzero(actual[num_valid:]).item() == 0 + + model._clear_deepstack_input_embeds(num_padded) + assert model.deepstack_input_embeds_num_tokens == 0 + for buffer in model.deepstack_input_embeds: + assert torch.count_nonzero(buffer[:num_padded]).item() == 0 + + decode_case = build_case(next_scenario, 1, 0, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + decode_result = _embed(model, decode_case) + lm_embed.assert_called_once() + assert_case_output(model, decode_case, decode_result) + decode_deepstack = model._get_deepstack_input_embeds(2) + assert decode_deepstack is not None + for level in range(_SHAPE.deepstack_levels): + assert torch.count_nonzero(decode_deepstack[f"deepstack_input_embeds_{level}"]).item() == 0 + + +@torch.inference_mode() +def test_deepstack_set_resizes_for_larger_scheduled_windows(mocker: MockerFixture): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16, buffer_capacity=4) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + for num_tokens in (16, 24): + case = build_case("mixed", num_tokens, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device="cpu") + lm_embed.reset_mock() + result = _embed(model, case) + + lm_embed.assert_called_once() + assert_case_output(model, case, result) + assert model.deepstack_input_embeds_num_tokens == num_tokens + assert all(buffer.size(0) >= num_tokens for buffer in model.deepstack_input_embeds) + model._clear_deepstack_input_embeds(num_tokens) + + +@pytest.mark.parametrize("mask_device", ["cpu", "cuda"]) +@pytest.mark.parametrize("scenario", ["image", "interleaved"]) +def test_cuda_embeddings_accept_cpu_or_cuda_mask(mocker: MockerFixture, scenario: str, mask_device: str): + model = make_thinker(_SHAPE, device="cuda", dtype=torch.bfloat16) + case = build_case(scenario, 16, 8, _SHAPE, device="cuda", dtype=torch.bfloat16, mask_device=mask_device) + lm_embed = mocker.spy(model.language_model, "embed_input_ids") + + result = _embed(model, case) + + lm_embed.assert_called_once() + assert result.device.type == "cuda" + assert_case_output(model, case, result) 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 c78679c55..5d5edf0df 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 @@ -1273,14 +1273,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - inputs_embeds = self._embed_text_input_ids( - input_ids, - self.language_model.embed_input_ids, - is_multimodal=is_multimodal, - ) - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds + return self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) # Detect interleaved audio-in-video early, since it affects # both the deepstack path and the final embedding merge. @@ -1295,12 +1293,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( is_interleaved = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio) - deepstack_input_embeds = None # split the feat dim to obtain multi-scale visual feature has_vision_embeddings = [ embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings ] - if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings): + has_deepstack_embeddings = self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings) + if has_deepstack_embeddings: multiscale_len = len(self.visual.deepstack_visual_indexes) multimodal_embeddings_multiscale = [] @@ -1343,6 +1341,28 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( if not is_interleaved: mm_position_idx += num_tokens + if is_interleaved: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + inputs_embeds = merge_interleaved_embeddings( + inputs_embeds, + multimodal_embeddings, + is_video, + is_audio, + is_mm_device, + ) + else: + # multimodal_embeddings now contains the main-scale features. + inputs_embeds = super().embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + if has_deepstack_embeddings: deepstack_input_embeds = inputs_embeds.new_zeros( inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1) ) @@ -1358,24 +1378,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) self._set_deepstack_input_embeds(deepstack_input_embeds) - if is_interleaved: - return merge_interleaved_embeddings( - inputs_embeds, - multimodal_embeddings, - is_video, - is_audio, - is_mm_device, - ) - - # Default: standard merge (no interleaving), same as parent class. - # multimodal_embeddings may have been updated above (deepstack - # main-scale). Use super() to stay consistent with the parent - # implementation and avoid issues seen in Qwen2.5-Omni (#34506). - return super().embed_input_ids( - input_ids, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) + return inputs_embeds def forward( self,