[Bugfix] Fix In Place Compile Overwrites for Qwen3Omni Embeddings with TP > 1 (#8532)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py

@@ -43,6 +43,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import (
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3MoeLLMModel,
     Qwen3OmniMoeThinkerForConditionalGeneration,
+    _get_capture_key,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -126,6 +127,29 @@ def _make_thinker(inner_output=None) -> tuple[Qwen3OmniMoeThinkerForConditionalG
     return thinker, inner
 
 
+def _overwrite_inputs_embeds(input_ids, positions, intermediate_tensors, *, inputs_embeds, **kwargs):
+    """Mutates the input embeddings by adding 1; we use this to check mutated behavior for compile."""
+    inputs_embeds.add_(1)
+    return inputs_embeds, {"hidden_states": {"layers": {0: inputs_embeds}}}
+
+
+def _overwrite_inputs_embeds_on_first_pp_rank(input_ids, positions, intermediate_tensors, *, inputs_embeds, **kwargs):
+    """Same as ``_overwrite_inputs_embeds``, on a PP rank that sends its captures to the next rank."""
+    inputs_embeds.add_(1)
+    return IntermediateTensors(
+        {"hidden_states": inputs_embeds, "residual": inputs_embeds, _get_capture_key(0): inputs_embeds}
+    )
+
+
+@pytest.fixture
+def first_rank_thinker(monkeypatch) -> Qwen3OmniMoeThinkerForConditionalGeneration:
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True),
+    )
+    return _make_thinker()[0]
+
+
 def _forward_args():
     return {
         "input_ids": torch.zeros(_TOKENS, dtype=torch.long),
@@ -160,6 +184,30 @@ def test_thinker_forward_passes_through_intermediate_tensors():
     assert out is it
 
 
+def test_thinker_embedding_capture_survives_overwritten_inputs_embeds(first_rank_thinker):
+    """Ensure the layer-0 capture is the original embeddings, not the overwritten inputs_embeds."""
+    first_rank_thinker.language_model.model = _overwrite_inputs_embeds
+    embeds = torch.randn(_TOKENS, _HIDDEN)
+
+    _, captured = first_rank_thinker.forward(
+        **_forward_args(), inputs_embeds=embeds.clone(), capture_layer_indices=[0], return_hidden_states=True
+    )
+
+    assert torch.equal(captured["hidden_states"]["layers"][0], embeds)
+
+
+def test_thinker_sends_original_embeddings_to_next_pp_rank(first_rank_thinker):
+    """Ensure a non-last PP rank sends the original embeddings as its layer-0 capture."""
+    first_rank_thinker.language_model.model = _overwrite_inputs_embeds_on_first_pp_rank
+    embeds = torch.randn(_TOKENS, _HIDDEN)
+
+    out = first_rank_thinker.forward(
+        **_forward_args(), inputs_embeds=embeds.clone(), capture_layer_indices=[0], return_hidden_states=True
+    )
+
+    assert torch.equal(out[_get_capture_key(0)], embeds)
+
+
 class _ThinkerStub(nn.Module):
     """Records capture kwargs; mirrors the fixed thinker return contract."""
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -144,6 +144,14 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 PP_CAPTURE_PREFIX = "capture_"
+# Capture index of the word embeddings, i.e. the input to decoder layer 0.
+_EMBEDDING_LAYER_INDEX = 0
+
+
+def _get_capture_key(layer_idx: int) -> str:
+    """Return the IntermediateTensors key that carries a layer capture across PP ranks."""
+    return f"{PP_CAPTURE_PREFIX}{layer_idx}"
+
 
 _THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration"
 
@@ -574,7 +582,7 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
                     hs = captured_hidden_states.setdefault("hidden_states", {})
                     layers = hs.setdefault("layers", {})
                     # Receive buffers are reused on the next step; retain an independent snapshot.
-                    layers[layer_idx] = intermediate_tensors[f"{PP_CAPTURE_PREFIX}{layer_idx}"].clone()
+                    layers[layer_idx] = intermediate_tensors[_get_capture_key(layer_idx)].clone()
 
         for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
             layer_idx = layer_idx + self.start_layer
@@ -602,7 +610,7 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
             if captured_hidden_states:
                 tensors.update(
                     {
-                        f"{PP_CAPTURE_PREFIX}{index}": value
+                        _get_capture_key(index): value
                         for index, value in captured_hidden_states["hidden_states"]["layers"].items()
                     }
                 )
@@ -1405,6 +1413,19 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         else:
             deepstack_input_embeds = None
 
+        # HACK: Snapshot the embedding capture here, before calling the compiled language model.
+        # This is needed for TP > 1, because vLLM's compile integration currently removes the
+        # clone on the graph input since it's read only, and enables kernels which write the
+        # residual in place, thereby causing the audio output to become corrupted.
+        embedding_capture = None
+        if (
+            return_hidden_states
+            and inputs_embeds is not None
+            and _EMBEDDING_LAYER_INDEX in (capture_layer_indices or ())
+        ):
+            embedding_capture = inputs_embeds.clone()
+            capture_layer_indices = [index for index in capture_layer_indices if index != _EMBEDDING_LAYER_INDEX]
+
         model_output = self.language_model.model(
             input_ids,
             positions,
@@ -1419,8 +1440,15 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         if inputs_embeds is not None and get_pp_group().is_first_rank:
             self._clear_deepstack_input_embeds(inputs_embeds.size(0))
 
+        if embedding_capture is not None:
+            # If we externally captured the embeddings for correct compile behavior, we need to rewrite the correct
+            # value back to the corresponding location to ensure read-only clone removal doesn't corrupt values
+            if isinstance(model_output, IntermediateTensors):
+                model_output[_get_capture_key(_EMBEDDING_LAYER_INDEX)] = embedding_capture
+            else:
+                model_output[-1]["hidden_states"]["layers"][_EMBEDDING_LAYER_INDEX] = embedding_capture
+
         if isinstance(model_output, IntermediateTensors):
-            # Non-last PP rank: forward the intermediate tensors as-is.
             return model_output
         hidden_states, captured_hidden_states = model_output
         if capture_layer_indices is None and not return_hidden_states: