From 48b3969717984d51480f890959f9e93461ea9ad5 Mon Sep 17 00:00:00 2001 From: Dnoob <92987618+Dnoob@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:51:02 +0800 Subject: [PATCH 1/5] [Bugfix] Fix Ovis Image crash when guidance_scale is set without negative_prompt (#1956) Signed-off-by: Dnoob --- vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py index eb62b83b3..8de617dd9 100644 --- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py +++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py @@ -685,6 +685,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi negative_text_ids = None if do_classifier_free_guidance: + negative_prompt = negative_prompt if negative_prompt is not None else "" negative_prompt_embeds, negative_text_ids = self.encode_prompt( prompt=negative_prompt, prompt_embeds=negative_prompt_embeds, From f9268822cdbdb02a1ddf3ca94582fbbea7e32adc Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Thu, 19 Mar 2026 16:51:45 +0800 Subject: [PATCH 2/5] [Bugfix] fix helios video generate use cpu device (#1915) Signed-off-by: rongfu.leng --- vllm_omni/diffusion/models/helios/pipeline_helios.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 8112cfd96..abccbec1a 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -897,13 +897,15 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip _, ph, pw = patch_size block_size = ph * pw + device = generator.device if generator is not None else self.device + cov = torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma cov += torch.eye(block_size) * 1e-8 cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. - L = torch.linalg.cholesky(cov) + L = torch.linalg.cholesky(cov).to(device) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) - z = torch.randn(block_number, block_size, generator=generator) + z = torch.randn(block_number, block_size, generator=generator, device=device) noise = z @ L.T noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw) From ba5a66b15b3a7c64133769c13880c445d5fd835b Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Thu, 19 Mar 2026 03:55:16 -0500 Subject: [PATCH 3/5] [XPU] update bagel modeling to remove cuda hardcode, add xpu stage_config (#1931) Signed-off-by: Chendi Xue Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- examples/offline_inference/bagel/end2end.py | 14 ++- .../model_executor/models/bagel/bagel.py | 5 +- .../platforms/xpu/stage_configs/bagel.yaml | 86 +++++++++++++++++++ 3 files changed, 100 insertions(+), 5 deletions(-) create mode 100644 vllm_omni/platforms/xpu/stage_configs/bagel.yaml diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index b1f308709..ae384cf37 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -34,7 +34,14 @@ def parse_args(): help="Path to input image for img2img.", ) - # Omni runtime init args + parser.add_argument( + "--output", + type=str, + default=".", + help="Output directory to save images.", + ) + + # OmniLLM init args parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-sleep-seconds", type=int, default=20) parser.add_argument("--batch-timeout", type=int, default=5) @@ -70,6 +77,7 @@ def parse_args(): def main(): args = parse_args() + os.makedirs(args.output, exist_ok=True) model_name = args.model prompts: list[OmniPromptType] = [] try: @@ -172,13 +180,13 @@ def main(): img_idx = 0 for req_output in omni_outputs: images = getattr(req_output, "images", None) + if not images: continue for j, img in enumerate(images): - save_path = f"output_{img_idx}_{j}.png" + save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") img.save(save_path) - print(f"[Info] Saved image to {save_path}") img_idx += 1 print(omni_outputs) diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index a053e98dd..e58b3501c 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -41,6 +41,7 @@ from vllm.multimodal.processing import ( ) from vllm.transformers_utils.processors.bagel import BagelProcessor +from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.bagel.autoencoder import ( AutoEncoderParams, DiagonalGaussian, @@ -439,7 +440,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): self._end_of_image_id = int(_tok.convert_tokens_to_ids("<|vision_end|>")) self._vae_token_mask: torch.Tensor | None = None - + self.device = get_local_device() self._install_mot_modules(config) def _install_mot_modules(self, config): @@ -627,7 +628,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): ) pos_embed = self.latent_pos_embed([vae_position_ids]) packed_timesteps = torch.tensor([timestep], device=padded_latent.device) - with torch.amp.autocast("cuda", dtype=torch.bfloat16): + with torch.amp.autocast(self.device.type, dtype=torch.bfloat16): timestep_embeds = self.time_embedder(packed_timesteps.to(padded_latent)) vae_embeds = self.vae2llm(latent) + timestep_embeds + pos_embed diff --git a/vllm_omni/platforms/xpu/stage_configs/bagel.yaml b/vllm_omni/platforms/xpu/stage_configs/bagel.yaml new file mode 100644 index 000000000..b1443d884 --- /dev/null +++ b/vllm_omni/platforms/xpu/stage_configs/bagel.yaml @@ -0,0 +1,86 @@ +# stage config for running bagel-7b-mot with architecture of OmniLLM. + +stage_args: + - stage_id: 0 + stage_type: llm + prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts + runtime: + devices: "0" + # 3 = 1 user prompt + 2 CFG companions (text-unconditional + image-unconditional). + max_batch_size: 3 + engine_args: + model_stage: thinker + model_arch: OmniBagelForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.9 + enforce_eager: true + trust_remote_code: true + engine_output_type: text + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 16384 + tensor_parallel_size: 1 + quantization: fp8 + omni_kv_config: + need_send_cache: true + kv_transfer_criteria: + type: prefill_finished #or special token generated + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 52 + detokenize: True + repetition_penalty: 1.05 + + - stage_id: 1 + stage_type: diffusion + cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches + runtime: + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: dit + gpu_memory_utilization: 0.9 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 1 + omni_kv_config: + need_recv_cache: true + engine_input_source: [0] + + final_output: true + final_output_type: image + is_comprehension: false + default_sampling_params: + seed: 52 + +# Runtime edges +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + # Distributed connectors configuration (optional) + # More connectors will be supported in the future. + connectors: + shared_memory_connector: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 # 64KB threshold + + + edges: + - from: 0 + to: 1 + window_size: -1 From 89fff09628208e924d723ebd00146b28edc37ed3 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Thu, 19 Mar 2026 17:10:08 +0800 Subject: [PATCH 4/5] [Fix] Fix slow hasattr in CUDAGraphWrapper.__getattr__ (#1982) Signed-off-by: ZeldaHuang --- tests/worker/test_cudagraph_wrapper_perf.py | 185 ++++++++++++++++++++ vllm_omni/worker/gpu_model_runner.py | 19 +- 2 files changed, 203 insertions(+), 1 deletion(-) create mode 100644 tests/worker/test_cudagraph_wrapper_perf.py diff --git a/tests/worker/test_cudagraph_wrapper_perf.py b/tests/worker/test_cudagraph_wrapper_perf.py new file mode 100644 index 000000000..d73fe46c9 --- /dev/null +++ b/tests/worker/test_cudagraph_wrapper_perf.py @@ -0,0 +1,185 @@ +"""Tests for CUDAGraphWrapper.__getattr__ performance optimization. + +This module tests that the patched CUDAGraphWrapper avoids expensive __repr__ +calls when hasattr() is used for non-existent attributes. The original vLLM +implementation includes {self.runnable} in the AttributeError message, which +triggers model tree traversal and can take ~6ms on large models. +""" + +import time + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class SlowReprModel(nn.Module): + """A mock model with artificially slow __repr__ to detect unwanted calls.""" + + def __init__(self, repr_delay_ms: float = 10.0): + super().__init__() + self.linear = nn.Linear(16, 16) + self.repr_delay_ms = repr_delay_ms + self.repr_call_count = 0 + + def forward(self, x): + return self.linear(x) + + def __repr__(self): + self.repr_call_count += 1 + # Simulate expensive repr by sleeping + time.sleep(self.repr_delay_ms / 1000.0) + return f"SlowReprModel(delay={self.repr_delay_ms}ms)" + + +class MockCUDAGraphWrapper: + """A minimal mock that mimics CUDAGraphWrapper structure for CPU testing.""" + + def __init__(self, runnable): + # Store in __dict__ directly to avoid triggering __getattr__ + object.__setattr__(self, "runnable", runnable) + + def __getattr__(self, key: str): + # This is the optimized implementation we're testing + runnable = object.__getattribute__(self, "runnable") + if hasattr(runnable, key): + return getattr(runnable, key) + # Key optimization: DO NOT include {self.runnable} in error message + # as it triggers expensive __repr__ on large models + raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper") + + +def test_hasattr_nonexistent_does_not_trigger_repr(): + """Verify that hasattr for non-existent attributes doesn't call __repr__.""" + model = SlowReprModel(repr_delay_ms=100.0) # Very slow repr + wrapper = MockCUDAGraphWrapper(model) + + # Reset counter + model.repr_call_count = 0 + + # Call hasattr for non-existent attribute multiple times + for _ in range(10): + result = hasattr(wrapper, "nonexistent_attribute_xyz") + assert result is False + + # __repr__ should never have been called + assert model.repr_call_count == 0, ( + f"__repr__ was called {model.repr_call_count} times when checking " + "for non-existent attributes. This indicates the AttributeError " + "message contains {self.runnable} which triggers expensive repr." + ) + + +def test_hasattr_nonexistent_is_fast(): + """Verify that hasattr for non-existent attributes is fast (<1ms per call).""" + model = SlowReprModel(repr_delay_ms=100.0) + wrapper = MockCUDAGraphWrapper(model) + + num_iterations = 100 + start = time.perf_counter() + for _ in range(num_iterations): + hasattr(wrapper, "nonexistent_attribute_xyz") + elapsed_ms = (time.perf_counter() - start) * 1000 + + avg_ms = elapsed_ms / num_iterations + # If __repr__ were being called, each would take ~100ms + # We expect <1ms per call with the fix + assert avg_ms < 1.0, ( + f"hasattr for non-existent attribute took {avg_ms:.2f}ms on average. " + "Expected <1ms. This suggests __repr__ is being triggered." + ) + + +def test_hasattr_existing_attribute_works(): + """Verify that hasattr for existing attributes returns True and works correctly.""" + model = SlowReprModel() + wrapper = MockCUDAGraphWrapper(model) + + # 'forward' exists on nn.Module + assert hasattr(wrapper, "forward") is True + + # 'linear' exists on our model + assert hasattr(wrapper, "linear") is True + + # Can actually access the attribute + linear = wrapper.linear + assert isinstance(linear, nn.Linear) + + +def test_getattr_existing_attribute_returns_value(): + """Verify that getattr for existing attributes returns the correct value.""" + model = SlowReprModel() + wrapper = MockCUDAGraphWrapper(model) + + # Access forward method + forward_method = wrapper.forward + assert callable(forward_method) + + # Access linear layer + linear = wrapper.linear + assert isinstance(linear, nn.Linear) + assert linear.in_features == 16 + assert linear.out_features == 16 + + +def test_getattr_nonexistent_raises_attribute_error(): + """Verify that getattr for non-existent attributes raises AttributeError.""" + model = SlowReprModel() + wrapper = MockCUDAGraphWrapper(model) + + with pytest.raises(AttributeError) as exc_info: + _ = wrapper.nonexistent_attribute + + # Verify error message format (should NOT contain model repr) + error_msg = str(exc_info.value) + assert "nonexistent_attribute" in error_msg + assert "cudagraph wrapper" in error_msg + # Should NOT contain the slow repr output + assert "SlowReprModel(delay=" not in error_msg + + +def test_attribute_error_message_does_not_contain_runnable_repr(): + """Explicitly verify the error message doesn't trigger runnable repr.""" + model = SlowReprModel(repr_delay_ms=100.0) + wrapper = MockCUDAGraphWrapper(model) + model.repr_call_count = 0 + + try: + _ = wrapper.nonexistent_attr + except AttributeError: + pass + + # __repr__ should not have been called during error construction + assert model.repr_call_count == 0, ( + "AttributeError message construction triggered __repr__. The error message should not include {self.runnable}." + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_real_cudagraph_wrapper_hasattr_performance(): + """Test the actual CUDAGraphWrapper from vllm_omni (requires CUDA).""" + from vllm.config import CUDAGraphMode + + model = SlowReprModel(repr_delay_ms=50.0).cuda() + model.repr_call_count = 0 + + # Create actual CUDAGraphWrapper + try: + wrapper = CUDAGraphWrapper(model, runtime_mode=CUDAGraphMode.NONE) + except Exception: + pytest.skip("Could not create CUDAGraphWrapper") + + # Test hasattr performance + num_iterations = 50 + start = time.perf_counter() + for _ in range(num_iterations): + hasattr(wrapper, "nonexistent_xyz") + elapsed_ms = (time.perf_counter() - start) * 1000 + + avg_ms = elapsed_ms / num_iterations + assert avg_ms < 1.0, f"Real CUDAGraphWrapper hasattr took {avg_ms:.2f}ms avg. Expected <1ms with the optimization." + assert model.repr_call_count == 0, f"__repr__ called {model.repr_call_count} times" diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 17acbe8b0..a536517f0 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,8 +1,9 @@ +import sys from typing import TYPE_CHECKING, Any, cast import numpy as np import torch -from vllm.compilation.cuda_graph import CUDAGraphWrapper +from vllm.compilation.cuda_graph import CUDAGraphWrapper as _OriginalCUDAGraphWrapper from vllm.config import CUDAGraphMode from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context @@ -36,6 +37,22 @@ else: logger = init_logger(__name__) +class CUDAGraphWrapper(_OriginalCUDAGraphWrapper): + def __getattr__(self, key: str) -> Any: + # allow accessing the attributes of the runnable. + if hasattr(self.runnable, key): + return getattr(self.runnable, key) + raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper") + + +# Patch vLLM's CUDAGraphWrapper with our optimized version +for _module_name, _module in sys.modules.items(): + if "vllm" not in _module_name: + continue + if hasattr(_module, "CUDAGraphWrapper") and _module.CUDAGraphWrapper is _OriginalCUDAGraphWrapper: + _module.CUDAGraphWrapper = CUDAGraphWrapper + + class OmniGPUModelRunner(GPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) From 60a8c01a968c46f26546620bb14789a118c8b72d Mon Sep 17 00:00:00 2001 From: rein yang <73573651+R2-Y@users.noreply.github.com> Date: Thu, 19 Mar 2026 17:34:01 +0800 Subject: [PATCH 5/5] [Bugfix] revert PR#1758 which introduced the accuracy problem of qwen3-omni (#2009) Signed-off-by: Rein Yang --- .../models/qwen3_omni/qwen3_omni.py | 50 +- .../qwen3_omni_moe_code_predictor_mtp.py | 548 ++++++++++-------- .../qwen3_omni/qwen3_omni_moe_talker.py | 113 ++-- vllm_omni/worker/gpu_model_runner.py | 9 +- 4 files changed, 424 insertions(+), 296 deletions(-) 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 6dcd278ac..ec928195a 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -159,12 +159,6 @@ class Qwen3OmniMoeForConditionalGeneration( # suppress tokens by setting their probability to ~1e-9 (finite very small) self.suppressed_tokens = self._get_talker_suppressed_tokens() self.requires_raw_input_tokens = True - # Keys that should stay on GPU in model_intermediate_buffer to avoid CPU↔GPU round-trips - self.gpu_resident_buffer_keys: set[str] = { - "last_talker_hidden", - "trailing_text_hidden", - "tts_pad_embed_projected", - } elif self.model_stage == "code2wav": self.enable_update_additional_information = True @@ -229,16 +223,14 @@ class Qwen3OmniMoeForConditionalGeneration( # ==================== Forward Pass ==================== def _get_talker_suppressed_tokens(self): - """Return a boolean mask on GPU for suppressed token positions.""" - vocab_size = self.config.talker_config.text_config.vocab_size - mask = torch.zeros(vocab_size, dtype=torch.bool) - start = vocab_size - 1024 - eos_id = self.config.talker_config.codec_eos_token_id - for i in range(start, vocab_size): - if i != eos_id: - mask[i] = True - # Will be moved to the correct device on first use - return mask + return [ + i + for i in range( + self.config.talker_config.text_config.vocab_size - 1024, + self.config.talker_config.text_config.vocab_size, + ) + if i != self.config.talker_config.codec_eos_token_id + ] def get_mrope_input_positions( self, @@ -586,7 +578,7 @@ class Qwen3OmniMoeForConditionalGeneration( Postprocess the talker hidden states. """ update_dict = {} - update_dict["last_talker_hidden"] = hidden_states[-1, :].detach() + update_dict["last_talker_hidden"] = hidden_states[-1, :].detach().to("cpu").contiguous() return update_dict def talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict): @@ -640,9 +632,9 @@ class Qwen3OmniMoeForConditionalGeneration( if inputs_embeds.shape[-1] == 2048: inputs_embeds = self.text_projection(inputs_embeds) code_predictor_codes, summed_embeddings = self.talker.code_predictor_forward( - input_ids, inputs_embeds, last_talker_hidden=last_talker_hidden + input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden ) - inputs_embeds = summed_embeddings + inputs_embeds = summed_embeddings.clone() inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size) return inputs_embeds, code_predictor_codes.squeeze(-1) @@ -766,7 +758,7 @@ class Qwen3OmniMoeForConditionalGeneration( # compatible with old shape [1,S,D] rem_tail = trailing_text_hidden.squeeze(0) if rem_tail.shape[0] > 0: - update_dict["trailing_text_hidden"] = rem_tail.detach() + update_dict["trailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous() # Also persist projected tts_pad for decode fallback if needed if isinstance(tts_pad_thinker, torch.Tensor): pad_in = tts_pad_thinker @@ -775,7 +767,7 @@ class Qwen3OmniMoeForConditionalGeneration( if pad_in.ndim == 1: pad_in = pad_in.view(1, 1, -1) pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker))) - update_dict["tts_pad_embed_projected"] = pad_proj.detach() + update_dict["tts_pad_embed_projected"] = pad_proj.detach().to("cpu").contiguous() except Exception: pass self._talker_cache_thinker_decode_embeds(info_dict, update_dict) @@ -934,7 +926,7 @@ class Qwen3OmniMoeForConditionalGeneration( if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: use_vec = q_tail[0:1, :] new_q_tail = ( - q_tail[1:, :].detach() + q_tail[1:, :].detach().to("cpu").contiguous() if q_tail.shape[0] > 1 else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) ) @@ -1130,10 +1122,16 @@ class Qwen3OmniMoeForConditionalGeneration( # implemented by assigning their logits to log(1e-9). if getattr(self, "model_stage", None) == "talker" and isinstance(logits, torch.Tensor): - # Move mask to device once (lazy), then reuse every step - if self.suppressed_tokens.device != logits.device: - self.suppressed_tokens = self.suppressed_tokens.to(logits.device) - logits.masked_fill_(self.suppressed_tokens.unsqueeze(0), -1e9) + try: + logits_cpu = logits.cpu() + logits_cpu[:, self.suppressed_tokens] = -1e9 + logits = logits_cpu.to(logits.device) + except Exception as e: + print(f"Error in logits suppression: {e}") + print(f"logits.shape: {logits.shape}") + print(f"self.suppressed_tokens: {self.suppressed_tokens}") + raise e + logits[:, self.suppressed_tokens] = -1e9 return logits def sample( diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index 2b823364f..fc7402890 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -1,16 +1,27 @@ -"""Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache. +"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support. -* SDPA attention (F.scaled_dot_product_attention) -- no HF backend fallback -* Persistent pre-allocated buffers (_proj_buf, _pos_ids) -- zero per-call alloc -* Inline top-k sampling -- no LogitsProcessorList / custom-op overhead -* torch.compile on inner transformer by default -* No @support_torch_compile / static_forward_context / namedtuple +This module implements the code predictor component for Qwen3-Omni talker models. + +The code predictor generates residual RVQ (Residual Vector Quantization) codes +autoregressively, predicting layers 1 to N based on layer-0 codes from the talker. """ +from collections import namedtuple +from typing import Any + import torch import torch.nn as nn import torch.nn.functional as F -from vllm.config import VllmConfig +from transformers import Cache, PretrainedConfig +from transformers.generation.logits_process import ( + LogitsProcessorList, + TopKLogitsWarper, + TopPLogitsWarper, +) +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -21,53 +32,47 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding - -from vllm_omni.platforms import current_omni_platform +from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) - -# =================================================================== -# Standalone Attention (SDPA, no KV cache, no HF backend fallback) -# =================================================================== +# ============================================================================ +# Code Predictor Attention Layer +# ============================================================================ class Qwen3OmniCodePredictorAttention(nn.Module): - """Multi-head self-attention for code predictor. - - Uses ``F.scaled_dot_product_attention`` directly. No KV cache -- the code - predictor always re-prefills the full (short) sequence each AR step. - - Input : [B, seq_len, hidden_size] - Output: [B, seq_len, hidden_size] - """ + """Multi-head self-attention for code predictor with vLLM optimization.""" def __init__( self, config, + layer_idx: int, + vllm_config: VllmConfig = None, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() - cp_cfg = config.code_predictor_config - self.num_heads = cp_cfg.num_attention_heads - self.num_kv_heads = cp_cfg.num_key_value_heads + + self.num_heads = config.code_predictor_config.num_attention_heads + self.num_key_value_heads = config.code_predictor_config.num_key_value_heads self.head_dim = getattr( - cp_cfg, + config.code_predictor_config, "head_dim", - cp_cfg.hidden_size // cp_cfg.num_attention_heads, + config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads, ) - self.hidden_size = cp_cfg.hidden_size - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads + self.hidden_size = config.code_predictor_config.hidden_size + + if self.num_heads % self.num_key_value_heads != 0: + raise ValueError("num_attention_heads must be divisible by num_key_value_heads") + + self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.qkv_proj = QKVParallelLinear( hidden_size=self.hidden_size, head_size=self.head_dim, total_num_heads=self.num_heads, - total_num_kv_heads=self.num_kv_heads, + total_num_kv_heads=self.num_key_value_heads, bias=False, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", @@ -83,55 +88,120 @@ class Qwen3OmniCodePredictorAttention(nn.Module): ) self.rotary_emb = get_rope( self.head_dim, - max_position=cp_cfg.max_position_embeddings, + max_position=config.code_predictor_config.max_position_embeddings, rope_parameters=None, dual_chunk_attention_config=None, ) - self.q_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_key_value_heads * self.head_dim + + # Query/Key normalization + self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) + self.is_causal = True + self.config = config + + self.attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"] + cudagraph_mode = get_current_vllm_config().compilation_config.cudagraph_mode + if "flash_attention_2" in ALL_ATTENTION_FUNCTIONS and cudagraph_mode.has_full_cudagraphs(): + logger.warning( + f"CUDAGraphMode.{cudagraph_mode.name} is currently not supported " + f"with flash attention for Qwen3-Omni talker MTP." + f"removing flash attention from attention_backends" + ) + self.attention_backends.remove("flash_attention_2") def forward( self, hidden_states: torch.Tensor, - position_ids: torch.Tensor, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool = False, + position_ids: torch.LongTensor | None = None, ) -> torch.Tensor: bsz, seq_len, _ = hidden_states.shape - # Flatten to 2-D so vLLM rotary_emb gets [num_tokens, size] - qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1)) + qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - # QK-norm -> RoPE (both 2-D) - q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) - k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + # Reshape for attention + q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim) + k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) + v = v.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) + + # Apply normalization + q = self.q_norm(q).contiguous() + k = self.k_norm(k).contiguous() + q = q.reshape(-1, self.q_size) + k = k.reshape(-1, self.kv_size) + + # Apply RoPE q, k = self.rotary_emb(position_ids, q, k) - # [B, heads, seq, head_dim] for SDPA - q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) - k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + # Reshape for attention + q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim) + k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) + v_heads = v.transpose(1, 2).contiguous() + q_heads = q.transpose(1, 2).contiguous() + k_heads = k.transpose(1, 2).contiguous() - attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1) - output, _ = self.o_proj(attn_out) - return output.view(bsz, seq_len, -1) + if past_key_values is not None: + sin, cos = self.rotary_emb.get_cos_sin(seq_len) + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs) + + # Try attention backends in order of preference, with runtime error handling + # This handles cases where the backend is registered but not actually available + attn_output = None + last_error = None + + for backend_name in self.attention_backends: + if backend_name not in ALL_ATTENTION_FUNCTIONS: + continue + + try: + attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name] + attn_output, _ = attention_interface( + self, + q_heads, + k_heads, + v_heads, + None, + dropout=0.0 if not self.training else getattr(self, "attention_dropout", 0.0), + scaling=self.head_dim**-0.5, + sliding_window=None, + use_cache=use_cache, + position_ids=position_ids[:seq_len].unsqueeze(0), + output_hidden_states=True, + output_attentions=False, + ) + break + except (ValueError, ImportError, RuntimeError, AttributeError) as e: + # Store error and try next backend + last_error = e + continue + + if attn_output is None: + raise RuntimeError( + f"All attention backends failed. Last error: {last_error}. " + "Please install flash-attn, or ensure PyTorch's scaled_dot_product_attention is available." + ) + attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous() + + attn_output, _ = self.o_proj(attn_output) + return attn_output -# =================================================================== -# MLP -# =================================================================== +# ============================================================================ +# Code Predictor MLP Layer +# ============================================================================ class Qwen3OmniCodePredictorMLP(nn.Module): - """SiLU-gated MLP for code predictor.""" + """Feed-forward network for code predictor with fused gate/up projection.""" def __init__( self, @@ -151,6 +221,7 @@ class Qwen3OmniCodePredictorMLP(nn.Module): prefix=f"{prefix}.gate_up_proj", disable_tp=True, ) + self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, @@ -167,23 +238,35 @@ class Qwen3OmniCodePredictorMLP(nn.Module): return down -# =================================================================== -# Decoder Layer -# =================================================================== +# ============================================================================ +# MTP Layer (Multi-Token Prediction Layer) +# ============================================================================ -class Qwen3OmniCodePredictorDecoderLayer(nn.Module): - """Transformer decoder layer (SDPA, no KV cache).""" +class Qwen3OmniCodePredictorMTPLayer(nn.Module): + """MTP layer for speculative decoding - predicts next residual code layer.""" def __init__( self, - config, + config: PretrainedConfig, + prefix: str, + model_config: ModelConfig, + layer_idx: int, + cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, - prefix: str = "", ) -> None: super().__init__() + self.layer_idx = layer_idx + self.config = config + self.self_attn = Qwen3OmniCodePredictorAttention( config, + layer_idx, + vllm_config=type( + "VllmConfig", + (), + {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config}, + )(), quant_config=quant_config, prefix=f"{prefix}.self_attn", ) @@ -192,247 +275,260 @@ class Qwen3OmniCodePredictorDecoderLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) - cp_cfg = config.code_predictor_config - self.input_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) - self.post_attention_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) + self.input_layernorm = RMSNorm( + config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps + ) + self.post_attention_layernorm = RMSNorm( + config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps + ) - def forward( + def mtp_block( self, hidden_states: torch.Tensor, - position_ids: torch.Tensor, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool = False, + position_ids: torch.LongTensor | None = None, ) -> torch.Tensor: + # Self-attention with residual residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = self.self_attn(hidden_states, past_key_values, cache_position, use_cache, position_ids) hidden_states = residual + hidden_states + # MLP with residual residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states + return hidden_states -# =================================================================== -# Base Transformer Model (re-prefill, no KV cache) -# =================================================================== - - class Qwen3OmniCodePredictorBaseModel(nn.Module): - """Inner transformer for code predictor. + """ + Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure. - Signature: ``forward(inputs_embeds, position_ids) -> hidden_states`` - -- plain Tensor in, plain Tensor out (no namedtuple). + This is a simple transformer that processes inputs_embeds and outputs hidden states. """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config.code_predictor_config - self.config = config + self.config = config + self.vocab_size = config.vocab_size + self.num_code_groups = config.num_code_groups + + # Codec embeddings (for layers 1-num_code_groups-1) self.codec_embedding = nn.ModuleList( - [VocabParallelEmbedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)] + [ + VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + for _ in range(config.num_code_groups - 1) + ] ) + # Decoder layers self.layers = nn.ModuleList( [ - Qwen3OmniCodePredictorDecoderLayer( + Qwen3OmniCodePredictorMTPLayer( vllm_config.model_config.hf_config, + f"{prefix}.layers.{idx}", + model_config=vllm_config.model_config, + layer_idx=idx, + cache_config=vllm_config.cache_config, quant_config=vllm_config.quant_config, - prefix=f"{prefix}.layers.{idx}", ) for idx in range(config.num_hidden_layers) ] ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, inputs_embeds: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Any | None = None, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + **kwargs: Any, + ) -> Any: + """ + Forward pass matching HF structure. + + Args: + inputs_embeds: [batch, seq_len, hidden_size] + position_ids: Optional position IDs tensor + past_key_values: Optional cached key-value pairs + use_cache: Whether to use cache + cache_position: Optional cache position tensor + **kwargs: Additional keyword arguments + + Returns: + Named tuple with .last_hidden_state and .past_key_values attributes + """ + batch_size, seq_len, _ = inputs_embeds.shape + # Forward through decoder layers hidden_states = inputs_embeds + for layer in self.layers: - hidden_states = layer(hidden_states, position_ids) + hidden_states = layer.mtp_block(hidden_states, past_key_values, cache_position, use_cache, position_ids) + + # Final norm hidden_states = self.norm(hidden_states) - return hidden_states + + # Return in HF-compatible format + Output = namedtuple("Output", ["last_hidden_state", "past_key_values"]) + return Output(last_hidden_state=hidden_states, past_key_values=None) # [batch, num_code_groups-1, hidden_size] + + def get_input_embeddings(self): + """Return codec embeddings for HF compatibility.""" + return self.codec_embedding -# =================================================================== -# Code Predictor Wrapper (optimized re-prefill, persistent buffers) -# =================================================================== +def code_predictor_sample( + logits: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + forward_context = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + logits = self.logits_processors(None, logits[:, -1]) + probs = F.softmax(logits, dim=-1) + code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1] + return code +def code_predictor_sample_fake( + logits: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + return torch.empty((logits.shape[0], 1), dtype=torch.int64, device=logits.device) + + +direct_register_custom_op( + op_name="qwen3_omni_code_predictor_sample", + op_func=code_predictor_sample, + fake_impl=code_predictor_sample_fake, +) + + +@support_torch_compile class Qwen3OmniMoeTalkerCodePredictor(nn.Module): - """Optimized code predictor -- re-prefill approach, no KV cache. + """ + Code predictor wrapper matching HF structure. - Each AR step forwards the full growing sequence (len 2 -> num_code_groups+1) - through the transformer. The extra O(T^2) FLOPs are negligible for - short sequences, and this avoids all KV-cache management overhead. - - Optimizations: - 1. Pre-allocated embedding buffer -- no torch.cat per step. - 2. Pre-allocated position_ids -- no torch.arange per step. - 3. Inline top-k sampling -- no LogitsProcessorList / custom op. - 4. Cached module references -- bypass ModuleList indexing. - 5. torch.compile on inner transformer. + Structure: + - self.model: Qwen3OmniCodePredictorBaseModel (transformer) + - self.lm_head: ModuleList of output heads """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - config = vllm_config.model_config.hf_config - self.config = config + talker_code_predictor_config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config self.prefix = prefix - self.num_code_groups = config.code_predictor_config.num_code_groups - self._hidden_size = config.code_predictor_config.hidden_size + self.config = talker_code_predictor_config + self.vocab_size = self.config.code_predictor_config.vocab_size + self.num_code_groups = self.config.code_predictor_config.num_code_groups - self.model = Qwen3OmniCodePredictorBaseModel( - vllm_config=vllm_config, - prefix=prefix, - ) + # Base transformer model (matches HF structure) + self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix) - # One lm_head per residual layer (layers 1 .. G-1) + # Output heads for each residual layer (1-num_layers-1) self.lm_head = nn.ModuleList( [ nn.Linear( - config.code_predictor_config.hidden_size, - config.code_predictor_config.vocab_size, + self.config.code_predictor_config.hidden_size, + self.config.code_predictor_config.vocab_size, bias=False, ) for _ in range(self.num_code_groups - 1) ] ) - - # Sampling hyperparams (inlined) - self._top_k = 50 - - # Persistent buffers (lazily initialised on first forward) - self._proj_buf: torch.Tensor | None = None - self._pos_ids: torch.Tensor | None = None - - # Cached plain-list refs (set once) - self._lm_heads: list | None = None - self._codec_embeds: list | None = None - - # Model forward (optionally compiled) - self._model_fwd: object | None = None - - # ------------------------------------------------------------------ - # Lazy-init helpers - # ------------------------------------------------------------------ - - def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: - max_seq = self.num_code_groups + 1 - if ( - self._proj_buf is not None - and self._proj_buf.shape[0] >= bsz - and self._proj_buf.device == device - and self._proj_buf.dtype == dtype - ): - return - self._proj_buf = torch.zeros(bsz, max_seq, self._hidden_size, dtype=dtype, device=device) - self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device) - - def _ensure_cached_refs(self) -> None: - if self._lm_heads is not None: - return - self._lm_heads = list(self.lm_head) - self._codec_embeds = list(self.model.codec_embedding) - - def _ensure_model_fwd(self) -> None: - if self._model_fwd is not None: - return - if not current_omni_platform.supports_torch_inductor(): - logger.warning_once("code_predictor: torch.compile disabled") - self._model_fwd = self.model.forward - return - self._model_fwd = torch.compile( - self.model.forward, - mode="default", - dynamic=True, + self.logits_processors = LogitsProcessorList( + [ + TopKLogitsWarper(top_k=50), + TopPLogitsWarper(top_p=0.8), + ] ) - logger.info("code_predictor: torch.compile enabled") - # ------------------------------------------------------------------ - # Forward -- re-prefill + persistent buffers + inline sampling - # ------------------------------------------------------------------ + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.layer_name = prefix - @torch.inference_mode() def forward( self, layer0_code: torch.Tensor, layer0_embed: torch.Tensor, last_talker_hidden: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Predict residual codebooks 1..G-1 autoregressively via re-prefill. + """ + Forward pass for code predictor. Args: - layer0_code: [bsz, 1] int64 - layer0_embed: [bsz, 1, hidden_size] - last_talker_hidden: [bsz, 1, hidden_size] + layer0_code: + Code index for code-group (layer) 0. + Shape: [batch_size, 1], dtype typically int64. + + last_talker_hidden: + + Shape: [batch_size, hidden_size]. Returns: - all_codes: [bsz, num_code_groups, 1] - proj_buf: [bsz, num_code_groups + 1, hidden_size] - pos 0 = last_talker_hidden (NOT a codec embed) - pos 1 = layer0_embed - pos 2.. = codec_embedding[i](predicted_code_i) + pos_all_layers: + Predicted codes for all code groups, including `layer0_code`. + Shape: [batch_size, num_code_groups, 1]. + + current_input: + The final input embedding sequence after appending embeddings of all + predicted codes (one token per predicted layer). + Shape: [batch_size, num_code_groups + 2, hidden_size]. """ - bsz = int(layer0_code.shape[0]) - device = layer0_code.device - dtype = last_talker_hidden.dtype - num_groups = self.num_code_groups - top_k = self._top_k + pos_codes = [layer0_code] # Start with layer 0: [batch, 1] + try: + current_input = torch.cat([last_talker_hidden, layer0_embed], dim=1) # [batch, 2, hidden_size] + except Exception as e: + print(f"Error in current_input: {e}") + print(f"last_talker_hidden shape: {last_talker_hidden.shape}") + print(f"prev_embed shape: {layer0_embed.shape}") + raise e + batch_size = current_input.shape[0] - # Lazy init - self._ensure_buffers(bsz, device, dtype) - self._ensure_model_fwd() - self._ensure_cached_refs() + # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively + for layer_idx in range(self.num_code_groups - 1): + seq_len = layer_idx + 2 + # Compute position_ids dynamically to avoid torch.compile specializing batch_size + position_ids = torch.arange(seq_len, device=current_input.device, dtype=torch.int64).repeat(batch_size) + # Forward through code_predictor model + outputs = self.model( + inputs_embeds=current_input, + attention_mask=None, + position_ids=position_ids, + past_key_values=None, + use_cache=False, + cache_position=None, + ) + hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size] - proj_buf = self._proj_buf - pos_ids = self._pos_ids - model_fwd = self._model_fwd - lm_heads = self._lm_heads - codec_embeds = self._codec_embeds - - # Output codes - all_codes = torch.empty(bsz, num_groups, 1, dtype=torch.int64, device=device) - all_codes[:, 0] = layer0_code - - # Fill buffer positions 0 & 1 - proj_buf[:bsz, 0:1, :] = last_talker_hidden - proj_buf[:bsz, 1:2, :] = layer0_embed - - # Autoregressive loop: predict layers 1..G-1 - for step in range(1, num_groups): - seq_len = step + 1 - projected = proj_buf[:bsz, :seq_len, :] - step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz) - - hidden_out = model_fwd(projected, step_pos_ids) - - # Inline top-k sampling - logits = lm_heads[step - 1](hidden_out[:, -1, :]) - if top_k > 0: - topk_vals, _ = logits.topk(top_k, dim=-1) - logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) - probs = F.softmax(logits, dim=-1) - code = torch.multinomial(probs, num_samples=1) # [bsz, 1] - - all_codes[:, step] = code - - # Embed predicted code -> next buffer position - new_embed = codec_embeds[step - 1](code) - proj_buf[:bsz, step + 1 : step + 2, :] = new_embed - - return all_codes, proj_buf[:bsz] - - # ------------------------------------------------------------------ - # Weight loading - # ------------------------------------------------------------------ + # Use the corresponding lm_head for this layer + logits = self.lm_head[layer_idx](hidden_state[:, -1:, :]) + code = torch.ops.vllm.qwen3_omni_code_predictor_sample(logits, self.layer_name) + pos_codes.append(code) + # Update prev_embed for next layer (if not last layer) + # layer_idx=0 predicts layer 1, embed with codec_embedding[1] + new_embed = self.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size] + current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size] + pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1] + return pos_all_layers, current_input def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with mapping for fused QKV and gate_up projections. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 9f5c20cf7..3362ae17f 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -129,70 +129,111 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") ) + max_batch_size = max( + vllm_config.scheduler_config.max_num_seqs, vllm_config.compilation_config.max_cudagraph_capture_size + ) + self.layer0_embed_buffer = torch.zeros( + (max_batch_size, 1, self.config.text_config.hidden_size), + dtype=vllm_config.model_config.dtype, + ) def code_predictor_forward( self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, *, + temperature: float = 1.0, + top_k: int = 50, # Match transformers default + top_p: float = 0.8, # Match transformers default + generation_steps: int | None = None, last_talker_hidden: torch.Tensor | None = None, **_: object, ) -> tuple[torch.Tensor, torch.Tensor]: - """Generate full RVQ codes + summed embeddings (single-loop, no KV cache). + """ + Generate full RVQ codec codes for the provided sequence. - The code predictor uses re-prefill: each AR step re-forwards the full - (short) sequence through the transformer. The returned ``proj_buf`` - already contains all codec embeddings at positions 1..G, - so summed_embeddings = proj_buf[:, 1:, :].sum(dim=1) — no second - loop or re-embedding needed. + The code predictor consumes the layer-0 codec codes produced by the talker + alongside the talker's hidden states, and autoregressively predicts the remaining + residual layers (to num_codec_groups). Returns: - result_codes: [batch, num_code_groups, seq_len] - summed_embeddings: [batch, seq_len, hidden_size] + tuple containing: + - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing + the complete set of codec codes + - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size] + Sum of all layer embeddings at each position (like Transformers) """ if input_ids is None: - raise ValueError("`input_ids` (layer-0 codes) must be provided.") + raise ValueError("`input_ids` containing layer-0 codec codes must be provided.") if inputs_embeds is None: - raise ValueError("`inputs_embeds` (talker hidden states) must be provided.") + raise ValueError("`inputs_embeds` containing talker hidden states must be provided.") if inputs_embeds.ndim == 2: inputs_embeds = inputs_embeds.unsqueeze(0) if input_ids.ndim == 1: input_ids = input_ids.unsqueeze(0) + # Ensure the tensors are contiguous for the autoregressive sampling loop + inputs_embeds = inputs_embeds.contiguous() + input_ids = input_ids.contiguous() + + # Generate full codec codes using MTP + # This will be the parallel prediction implementation batch_size, seq_len = input_ids.shape - device = input_ids.device - embed_fn = self.language_model.model.codec_embedding - hidden_size = self.config.code_predictor_config.hidden_size - result_codes = torch.empty( - batch_size, - self.num_code_groups, - seq_len, - dtype=torch.int64, - device=device, - ) - summed_embeddings = torch.empty( - batch_size, - seq_len, - hidden_size, - dtype=inputs_embeds.dtype, - device=device, - ) + # For now, use sequential generation (TODO: implement parallel) + # Result will be [batch, num_code_groups, seq_len] + # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position + all_codes_per_position = [] + middle_hidden_states = [] # Collect hidden states for each position + # Generate residual layers for each position for pos in range(seq_len): - layer0_code = input_ids[:, pos : pos + 1] - layer0_embed = embed_fn(layer0_code) + layer0_code = input_ids[:, pos : pos + 1] # [batch, 1] - pos_all_layers, proj_buf = self.code_predictor( - layer0_code, - layer0_embed, - last_talker_hidden, + # Initial input: [last_talker_hidden, layer0_embed] + layer0_embed = self.embed_input_ids(layer0_code) + self.layer0_embed_buffer[:batch_size].copy_(layer0_embed) + pos_all_layers, current_input = self.code_predictor( + layer0_code, self.layer0_embed_buffer[:batch_size], last_talker_hidden ) - result_codes[:, :, pos : pos + 1] = pos_all_layers - # proj_buf layout: [0]=talker_hidden, [1..G]=codec embeds - summed_embeddings[:, pos, :] = proj_buf[:, 1:, :].sum(dim=1) + # Stack all layers for this position: [batch, num_code_groups, 1] + all_codes_per_position.append(pos_all_layers) + middle_hidden_states.append(current_input[:, 2:-1, :]) + + # Concatenate across positions: [batch, num_code_groups, seq_len] + result_codes = torch.cat(all_codes_per_position, dim=2) + + # Build summed embeddings for each position (like Transformers) + # This combines layer-0 embed, mid layers hidden states, and last layer embed + all_summed_embeddings = [] + + for pos in range(seq_len): + # Layer 0 embedding + layer0_code = result_codes[:, 0, pos : pos + 1] # [batch, 1] + layer0_embed = self.embed_input_ids(layer0_code) # [batch, 1, hidden_size] + + # mid layers hidden states (from CodePredictor) + mid_residual_hiddens = middle_hidden_states[pos] # [batch, num_code_groups-2, hidden_size] + mid_list = list(mid_residual_hiddens.split(1, dim=1)) + + # last layer embedding + last_layer_code = result_codes[:, -1, pos : pos + 1] # [batch, 1] + last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code) + + # Concatenate all layers: [batch, num_code_groups, hidden_size] + pos_codec_hiddens = torch.cat( + [layer0_embed] + mid_list + [last_residual_hidden], + dim=1, + ) + + # Sum across layers: [batch, 1, hidden_size] (like Transformers) + pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True) + all_summed_embeddings.append(pos_summed) + + # Concatenate across positions: [batch, seq_len, hidden_size] + summed_embeddings = torch.cat(all_summed_embeddings, dim=1).squeeze(1) return result_codes, summed_embeddings diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a536517f0..af3006ab2 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1358,17 +1358,10 @@ class OmniGPUModelRunner(GPUModelRunner): req_state = self.requests.get(req_id) if req_state is None: return - # Check if the model declares keys that should stay on GPU - gpu_keys: set[str] = set() - if hasattr(self, "model") and hasattr(self.model, "gpu_resident_buffer_keys"): - gpu_keys = self.model.gpu_resident_buffer_keys existing = self.model_intermediate_buffer.setdefault(req_id, {}) for k, v in upd.items(): if isinstance(v, torch.Tensor): - if k in gpu_keys: - existing[k] = v.detach().clone() - else: - existing[k] = v.detach().to("cpu").contiguous() + existing[k] = v.detach().to("cpu").contiguous() elif isinstance(v, list): existing[k] = [ (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v