# 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)