[Qwen3-Omni] Remove redundant text embedding in DeepStack path (#7573)

Signed-off-by: howard-shan <shanwenhao668@gmail.com>
Signed-off-by: howard.shan <shanwenhao668@gmail.com>

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)

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)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -1273,15 +1273,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         *,
         is_multimodal: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        inputs_embeds = self._embed_text_input_ids(
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return 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
-
         # Detect interleaved audio-in-video early, since it affects
         # both the deepstack path and the final embedding merge.
         # is_multimodal from the worker may be CPU; bitwise ops need input_ids' device.
@@ -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,