From 63bbd19a1d04aa7138528b6ee37cf9e17fc28386 Mon Sep 17 00:00:00 2001 From: Khairul Kabir <250204754+khairulkabir1661@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:19:41 -0700 Subject: [PATCH] [Bugfix] Read hidden layer key from model config instead of hardcoding 24 (#5825) Signed-off-by: khairulkabir1661 Co-authored-by: khairulkabir1661 --- .../test_qwen3_omni_streaming_helpers.py | 38 ++++- .../stage_input_processors/qwen3_omni.py | 155 +++++++----------- 2 files changed, 95 insertions(+), 98 deletions(-) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 208cc72c6..dfbb9d91d 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -282,8 +282,14 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None: "hidden_states.layer_24": torch.full((3, 2), 2.0), "embed.tts_bos": torch.zeros(1, 2), } + # Mock transfer_manager with model config + transfer_manager = SimpleNamespace( + _get_model_config=lambda: SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=24)) + ) + ) - payload = q3.thinker2talker_full_payload(None, pooling_output, request) + payload = q3.thinker2talker_full_payload(transfer_manager, pooling_output, request) assert payload is not None assert payload["ids"]["all"] == [151644, 872, 3] @@ -293,6 +299,36 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None: assert payload["hidden_states"]["output"].shape[0] == 2 +def test_thinker2talker_full_payload_uses_config_hidden_layer() -> None: + """accept_hidden_layer=6 picks up layer_6, not the old hardcoded 24.""" + request = SimpleNamespace( + request_id="thinker", + prompt_token_ids=[151644, 872], + output_token_ids=[3], + all_token_ids=[151644, 872, 3], + ) + layer_6_data = torch.full((3, 2), 6.0) + layer_24_data = torch.full((3, 2), 24.0) + pooling_output = { + "hidden_states.layer_0": torch.ones(3, 2), + "hidden_states.layer_6": layer_6_data, + "hidden_states.layer_24": layer_24_data, + "embed.tts_bos": torch.zeros(1, 2), + } + transfer_manager = SimpleNamespace( + _get_model_config=lambda: SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=6)) + ) + ) + + payload = q3.thinker2talker_full_payload(transfer_manager, pooling_output, request) + + assert payload is not None + # Should use layer 6 (value 6.0), not layer 24 (value 24.0) + assert payload["hidden_states"]["output"].shape[0] == 2 + assert torch.allclose(payload["hidden_states"]["output"], layer_6_data[:-1]) + + def test_thinker2talker_token_only_preserves_voice_metadata() -> None: source_outputs = [ SimpleNamespace( diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 7e28b94c3..3cef0eed8 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -32,9 +32,8 @@ from vllm_omni.model_executor.stage_input_processors.tts_utils import ( logger = logging.getLogger(__name__) -# Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer -_EMBED_LAYER_KEY = "0" -_HIDDEN_LAYER_KEY = "24" +# Pooling output layer indices (semantic names for clarity) +# Actual layer indices are integers; these functions provide semantic access # Per-model REPLACE-keys for the full-payload accumulator. Keys in this # set use REPLACE semantics (subsequent emissions discard prior chunks) # instead of CONCAT. qwen3-omni currently has none — model_outputs is @@ -48,16 +47,35 @@ _QWEN3_CODEC_EOS_TOKEN_ID = 4198 def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: - """Fetch layer tensor with tolerant key lookup (str/int).""" + """Fetch layer tensor with tolerant key lookup (str/int/semantic).""" if not isinstance(layers, dict): return None - key_int = int(key) - val = layers.get(key_int) + # Try direct string key first (supports semantic keys like "embedding") + val = layers.get(key) if val is None: - val = layers.get(key) + # Try converting to int for backward compatibility with numeric string keys + try: + key_int = int(key) + val = layers.get(key_int) + except (ValueError, TypeError): + pass return val if isinstance(val, torch.Tensor) else None +def _get_embedding_layer_index() -> int: + """Return the layer index for word embeddings (always layer 0).""" + return 0 + + +def _get_accept_hidden_layer_index(transfer_manager: Any) -> int: + """Read accept_hidden_layer index from model config.""" + if hasattr(transfer_manager, "_get_model_config"): + mc = transfer_manager._get_model_config() + else: + mc = getattr(transfer_manager, "config", None) + return int(mc.hf_config.talker_config.accept_hidden_layer) + + def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: im_start_token_id = 151644 system_token_id = 8948 @@ -174,88 +192,6 @@ def _extract_qwen3_full_payload_codec_rows( } -# ========================= -# PD disaggregation helpers -# ========================= - - -def _get_prefill_multimodal_output( - request_id: str, - streaming_context: Any | None, -) -> dict[str, Any] | None: - bridge_states = getattr(streaming_context, "bridge_states", None) - if not isinstance(bridge_states, dict): - return None - by_req = bridge_states.get("pd_prefill_multimodal_output_by_req") - if not isinstance(by_req, dict): - return None - prefill_mm = by_req.get(request_id) - return prefill_mm if isinstance(prefill_mm, Mapping) else None - - -def _merge_pd_embeddings( - decode_emb: torch.Tensor, - decode_hid: torch.Tensor, - prefill_mm: dict[str, Any], - device: torch.device, - expected_total: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Merge prefill prompt embeddings with decode generated embeddings. - - In PD mode the prefill engine processes the prompt and the decode engine - generates tokens starting from position 1. This function concatenates - them, removing the overlapping token(s): - - merged = prefill[:P] + decode[overlap:] - - where overlap = P + D - expected_total. - """ - try: - p_layers = prefill_mm.get("hidden_states", {}).get("layers", {}) - p_emb = p_layers[int(_EMBED_LAYER_KEY)].detach().to(device=device, dtype=torch.float) - p_hid = p_layers[int(_HIDDEN_LAYER_KEY)].detach().to(device=device, dtype=torch.float) - except (KeyError, AttributeError, TypeError) as exc: - available_keys = list(prefill_mm.keys()) if isinstance(prefill_mm, Mapping) else type(prefill_mm).__name__ - logger.error( - "_merge_pd_embeddings: failed to extract prefill embeddings (%s). " - "Expected keys %r and %r, got: %s. " - "Falling back to decode-only embeddings – talker user-segment will be degraded.", - exc, - _EMBED_LAYER_KEY, - _HIDDEN_LAYER_KEY, - available_keys, - ) - return decode_emb, decode_hid - - if p_emb.shape[0] == 0 or decode_emb.shape[0] == 0: - return decode_emb, decode_hid - - raw_total = p_emb.shape[0] + decode_emb.shape[0] - overlap = max(0, raw_total - expected_total) if expected_total is not None else 0 - - merged_emb = torch.cat([p_emb, decode_emb[overlap:]], dim=0) - merged_hid = torch.cat([p_hid, decode_hid[overlap:]], dim=0) - return merged_emb, merged_hid - - -def _resolve_tts_token_embedding( - key: str, - *, - thinker_mm: dict[str, Any], - prefill_mm: dict[str, Any] | None, - device: torch.device, -) -> torch.Tensor | None: - """Return TTS BOS/EOS/PAD embedding tensors for the talker projection path. - - Values are taken from the current thinker (decode) ``multimodal_output``; in - PD mode, missing keys may be filled from the paired prefill stage output. - """ - val = thinker_mm.get("embed", {}).get(key) - if val is None and prefill_mm is not None: - val = prefill_mm.get("embed", {}).get(key) - return val.detach().to(device=device, dtype=torch.float) if val is not None else None - - # ========================= # Streaming input helpers # ========================= @@ -454,13 +390,32 @@ def thinker2talker_async_chunk( thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {} thinker_embed_raw = multimodal_output.get("embed", {}) thinker_embed = thinker_embed_raw if isinstance(thinker_embed_raw, dict) else {} - thinker_emb = _layer_tensor(thinker_layers, _EMBED_LAYER_KEY) - thinker_hid = _layer_tensor(thinker_layers, _HIDDEN_LAYER_KEY) + + embedding_idx = _get_embedding_layer_index() + accept_hidden_idx = _get_accept_hidden_layer_index(transfer_manager) + + # Remap numeric layer indices to semantic keys; use explicit None checks + # because multi-element tensors raise on boolean evaluation. + emb_val = thinker_layers.get(embedding_idx) + if emb_val is None: + emb_val = thinker_layers.get(str(embedding_idx)) + hid_val = thinker_layers.get(accept_hidden_idx) + if hid_val is None: + hid_val = thinker_layers.get(str(accept_hidden_idx)) + semantic_layers = { + "embedding": emb_val, + "accept_hidden": hid_val, + } + thinker_emb = _layer_tensor(semantic_layers, "embedding") + thinker_hid = _layer_tensor(semantic_layers, "accept_hidden") if thinker_emb is None or thinker_hid is None: logger.debug( - "thinker2talker_async_chunk: missing thinker layers for req=%s (embed=%s hidden=%s)", + "thinker2talker_async_chunk: missing thinker layers for req=%s " + "(embedding_layer=%d:%s accept_hidden_layer=%d:%s)", request_id, + embedding_idx, thinker_emb is not None, + accept_hidden_idx, thinker_hid is not None, ) return None @@ -542,21 +497,27 @@ def thinker2talker_full_payload( ) return None + embedding_idx = _get_embedding_layer_index() + accept_hidden_idx = _get_accept_hidden_layer_index(transfer_manager) + + # Use semantic keys instead of numeric indices layers = { - 0: pooling_output.get("hidden_states.layer_0"), - 24: pooling_output.get("hidden_states.layer_24"), + "embedding": pooling_output.get(f"hidden_states.layer_{embedding_idx}"), + "accept_hidden": pooling_output.get(f"hidden_states.layer_{accept_hidden_idx}"), } - thinker_emb = _layer_tensor(layers, _EMBED_LAYER_KEY) - thinker_hid = _layer_tensor(layers, _HIDDEN_LAYER_KEY) + thinker_emb = _layer_tensor(layers, "embedding") + thinker_hid = _layer_tensor(layers, "accept_hidden") if thinker_emb is None: hidden = pooling_output.get("hidden") thinker_emb = hidden if isinstance(hidden, torch.Tensor) else None if thinker_emb is None or thinker_hid is None: logger.warning( "thinker2talker_full_payload: missing thinker tensors for req=%s " - "(embed=%s hidden=%s keys=%s); consumer wait gate may hang.", + "(embedding_layer=%d:%s accept_hidden_layer=%d:%s keys=%s); consumer wait gate may hang.", rid, + embedding_idx, thinker_emb is not None, + accept_hidden_idx, thinker_hid is not None, list(pooling_output.keys()), )