Merge remote-tracking branch 'remotes/origin/main' into dev/vllm-align

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)

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"

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)

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,

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
 

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(

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.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -128,70 +128,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
 

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

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)
@@ -1356,17 +1373,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