[Refactor] Standardize data entry key names to {type}.{qualifier} format (#1829)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -12,6 +12,7 @@ from pytest_mock import MockerFixture
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
@@ -111,7 +112,11 @@ def test_load_poll(build_adapter):
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.load_async(request)
-    payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True}
+    payload: OmniPayload = {
+        "codes": {"audio": [[1]]},
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 16)
     adapter._poll_single_request(request)
 
@@ -154,12 +159,49 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
-    adapter._update_request_payload("ext", {"h": torch.tensor([[1.0]]), "codes": [1], "finished": False})
-    merged = adapter._update_request_payload("ext", {"h": torch.tensor([[2.0]]), "codes": [2], "finished": True})
+    first: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "codes": {"audio": [1]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    adapter._update_request_payload("ext", first)
+    second: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "codes": {"audio": [2]},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    merged = adapter._update_request_payload("ext", second)
 
-    assert torch.equal(merged["h"], torch.tensor([[1.0], [2.0]]))
-    assert merged["codes"] == [1, 2]
-    assert merged["finished"] is True
+    assert torch.equal(merged["hidden_states"]["output"], torch.tensor([[1.0], [2.0]]))
+    assert merged["codes"]["audio"] == [1, 2]
+    assert merged["meta"]["finished"].item() is True
+
+
+def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
+
+    adapter.request_ids_mapping["req-merged"] = "ext-merged"
+    adapter.request_payload["ext-merged"] = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "ids": {"prompt": [11, 12]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    connector.get.return_value = (payload, 8)
+
+    adapter._poll_single_request(request)
+
+    assert torch.equal(
+        request.additional_information["hidden_states"]["output"],
+        torch.tensor([[1.0], [2.0]]),
+    )
+    # Keys absent from the new chunk are dropped (matches main's behavior).
+    assert "ids" not in request.additional_information
+    assert request.additional_information["meta"]["finished"].item() is True
 
 
 def test_process_and_restore_queues(build_adapter):
@@ -321,7 +363,10 @@ def test_cleanup_after_poll_flow(build_adapter):
     adapter.load_async(request)
 
     adapter.request_ids_mapping["req-flow"] = "ext-flow"
-    payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True}
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 8)
     adapter._poll_single_request(request)
 

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -44,7 +44,7 @@ def test_forward_trims_context_on_exact_frame_boundaries():
 
     out = model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
-        runtime_additional_information=[{"left_context_size": 2}],
+        runtime_additional_information=[{"meta": {"left_context_size": 2}}],
     )
 
     audio = out.multimodal_outputs["model_outputs"][0]
@@ -57,7 +57,7 @@ def test_forward_trims_trailing_padding_without_context():
 
     out = model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
-        runtime_additional_information=[{"left_context_size": 0}],
+        runtime_additional_information=[{"meta": {"left_context_size": 0}}],
     )
 
     audio = out.multimodal_outputs["model_outputs"][0]

tests/model_executor/models/test_fish_speech_regressions.py

@@ -45,7 +45,7 @@ def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
 
     out = decoder.forward(
         input_ids=torch.arange(20, dtype=torch.long),
-        runtime_additional_information=[{}, {"left_context_size": 1}],
+        runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}],
     )
 
     audios = out.multimodal_outputs["model_outputs"]

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -58,8 +58,8 @@ def test_text2flow_supports_batched_source_outputs():
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
     assert outputs[1]["prompt_token_ids"] == [4, 5]
-    assert outputs[0]["additional_information"]["prefix_ids"] == [10, 11]
-    assert outputs[1]["additional_information"]["prefix_ids"] == [20, 21]
+    assert outputs[0]["additional_information"]["ids"]["prompt"] == [10, 11]
+    assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
@@ -83,8 +83,8 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     )
 
     assert payload is not None
-    assert payload["finished"].item() is True
-    assert payload["code_predictor_codes"] == [1, 2, 3]
+    assert payload["meta"]["finished"].item() is True
+    assert payload["codes"]["audio"] == [1, 2, 3]
     assert payload["token_offset"] == 0
     assert payload["left_context_size"] == 0
     assert payload["req_id"] == ["rid-0"]
@@ -111,8 +111,8 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes
     )
 
     assert payload is not None
-    assert payload["code_predictor_codes"] == []
-    assert payload["finished"].item() is True
+    assert payload["codes"]["audio"] == []
+    assert payload["meta"]["finished"].item() is True
 
 
 def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
@@ -138,7 +138,7 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
     )
 
     assert payload1 is not None
-    assert payload1["code_predictor_codes"] == [1, 2]
+    assert payload1["codes"]["audio"] == [1, 2]
     assert payload1["token_offset"] == 0
     assert payload2 is None
 
@@ -168,9 +168,9 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["code_predictor_codes"] == [1, 2, 3]
+    assert payload_ready["codes"]["audio"] == [1, 2, 3]
     assert payload_ready["token_offset"] == 0
-    assert payload_ready["finished"].item() is False
+    assert payload_ready["meta"]["finished"].item() is False
 
 
 def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
@@ -197,12 +197,12 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
     )
 
     assert payload_stream is not None
-    assert payload_stream["finished"].item() is False
-    assert payload_stream["code_predictor_codes"] == [3, 4, 5]
+    assert payload_stream["meta"]["finished"].item() is False
+    assert payload_stream["codes"]["audio"] == [3, 4, 5]
     assert payload_stream["token_offset"] == 0
     assert payload_final is not None
-    assert payload_final["finished"].item() is True
-    assert payload_final["code_predictor_codes"] == [3, 4, 5, 6]
+    assert payload_final["meta"]["finished"].item() is True
+    assert payload_final["codes"]["audio"] == [3, 4, 5, 6]
     assert payload_final["token_offset"] == 2
 
 
@@ -233,7 +233,7 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["code_predictor_codes"] == [8, 9, 10, 11]
+    assert payload_ready["codes"]["audio"] == [8, 9, 10, 11]
     assert payload_ready["token_offset"] == 0
 
 
@@ -260,8 +260,8 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     )
 
     assert payload_stream is not None
-    assert payload_stream["finished"].item() is False
-    assert payload_stream["code_predictor_codes"] == [3, 4]
+    assert payload_stream["meta"]["finished"].item() is False
+    assert payload_stream["codes"]["audio"] == [3, 4]
     assert payload_final is not None
-    assert payload_final["finished"].item() is True
-    assert payload_final["code_predictor_codes"] == []
+    assert payload_final["meta"]["finished"].item() is True
+    assert payload_final["codes"]["audio"] == []

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -299,7 +299,7 @@ class TestAr2Diffusion:
     def test_i2i_with_mm_output(self):
         """Test image-to-image with prior_token_image_ids from AR model."""
         token_ids = list(range(1024)) + [16385]
-        mm_output = {"prior_token_image_ids": torch.tensor([1, 2, 3])}
+        mm_output = {"ids": {"prior_image": torch.tensor([1, 2, 3])}}
         stage_list = [_stage_with_outputs([_source_output(token_ids, mm_output)])]
 
         from PIL import Image

tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py

@@ -13,23 +13,23 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _sentinel():
-    return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+    return {"codes": {"audio": []}, "meta": {"finished": torch.tensor(True, dtype=torch.bool)}}
 
 
 def test_flush_remaining_codes_when_no_codes_accumulated_missing_request_id():
     """No entry for request_id: treat as empty, return finished sentinel with empty codes."""
     tm = SimpleNamespace(code_prompt_token_ids={})
     out = _flush_remaining_codes(tm, "missing", chunk_size=3, left_context_size=3)
-    assert out["code_predictor_codes"] == _sentinel()["code_predictor_codes"]
-    assert out["finished"].equal(_sentinel()["finished"])
+    assert out["codes"]["audio"] == _sentinel()["codes"]["audio"]
+    assert out["meta"]["finished"].item() is True
 
 
 def test_flush_remaining_codes_when_no_codes_accumulated_empty_list():
     """Explicit empty accumulation list returns the same sentinel."""
     tm = SimpleNamespace(code_prompt_token_ids={"r": []})
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["code_predictor_codes"] == []
-    assert out["finished"].item() is True
+    assert out["codes"]["audio"] == []
+    assert out["meta"]["finished"].item() is True
 
 
 def test_flush_remaining_codes_partial_chunk_remaining():
@@ -41,8 +41,8 @@ def test_flush_remaining_codes_partial_chunk_remaining():
         code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6], [7]]},
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["finished"].item() is True
-    assert out["code_predictor_codes"] == [4, 5, 6, 7]
+    assert out["meta"]["finished"].item() is True
+    assert out["codes"]["audio"] == [4, 5, 6, 7]
 
 
 def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
@@ -52,7 +52,7 @@ def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
     # context_length = chunk_size = 3, end_index = min(6, 6) -> all 6
-    assert out["code_predictor_codes"] == [1, 2, 3, 4, 5, 6]
+    assert out["codes"]["audio"] == [1, 2, 3, 4, 5, 6]
 
 
 @pytest.mark.parametrize(
@@ -74,5 +74,5 @@ def test_flush_remaining_codes_context_window_end_index(
     tm = SimpleNamespace(code_prompt_token_ids={"r": accumulated})
     out = _flush_remaining_codes(tm, "r", chunk_size=chunk_size, left_context_size=left_context)
     expected_flat = list(range(length - expected_end_index, length))
-    assert out["code_predictor_codes"] == expected_flat
-    assert out["finished"].item() is True
+    assert out["codes"]["audio"] == expected_flat
+    assert out["meta"]["finished"].item() is True

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -19,7 +19,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"):
     """Build a minimal stage_list[0] with a single talker output carrying codec_codes."""
-    output = SimpleNamespace(multimodal_output={"code_predictor_codes": codec_codes})
+    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
     talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
     stage0 = SimpleNamespace(engine_outputs=[talker_output])
     return [stage0]

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -55,7 +55,7 @@ def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
     return talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,))},
+        pooling_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req(rid, finished=finished, initial_codec_chunk_frames=req_ic),
         is_finished=finished,
     )
@@ -65,7 +65,7 @@ def test_empty_returns_none():
     tm = _tm()
     p = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,))},
+        pooling_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req("r", finished=False),
     )
     assert p is None
@@ -79,7 +79,8 @@ def test_eof_marker_when_finished_empty():
         request=_req("r", finished=True),
         is_finished=True,
     )
-    assert p == {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+    assert p["codes"] == {"audio": []}
+    assert p["meta"]["finished"].item() is True
 
 
 def test_flush_on_finish():
@@ -92,8 +93,8 @@ def test_flush_on_finish():
         is_finished=True,
     )
     assert p is not None
-    assert p["finished"] is True
-    assert len(p["code_predictor_codes"]) == _Q * 24
+    assert p["meta"]["finished"].item() is True
+    assert len(p["codes"]["audio"]) == _Q * 24
 
 
 _CASES = [
@@ -158,8 +159,8 @@ def test_streaming_phases(config, n_frames, finished, expected):
     else:
         exp_ctx, exp_window = expected
         assert payload is not None
-        assert payload["left_context_size"] == exp_ctx
-        assert len(payload["code_predictor_codes"]) == _Q * exp_window
+        assert payload["meta"]["left_context_size"] == exp_ctx
+        assert len(payload["codes"]["audio"]) == _Q * exp_window
 
 
 def test_dynamic_ic_adapts_to_load():
@@ -169,14 +170,14 @@ def test_dynamic_ic_adapts_to_load():
     # Low load (1/8) -> IC=2 -> emit at 2
     p1 = _call(tm, "r", n_frames=2)
     assert p1 is not None
-    assert len(p1["code_predictor_codes"]) == _Q * 2
+    assert len(p1["codes"]["audio"]) == _Q * 2
 
     # High load: add 4 others -> active=5/8 -> IC=8 -> emit at 8
     for i in range(4):
         tm.code_prompt_token_ids[f"other-{i}"] = [[0]]
     p2 = _call(tm, "r", n_frames=8)
     assert p2 is not None
-    assert len(p2["code_predictor_codes"]) == _Q * 8
+    assert len(p2["codes"]["audio"]) == _Q * 8
 
     # Requests past initial phase still count in load factor
     tm2 = _tm(max_num_seqs=4)
@@ -185,7 +186,7 @@ def test_dynamic_ic_adapts_to_load():
     # active=4/4=1.0 -> IC=16
     p3 = _call(tm2, "new", n_frames=16)
     assert p3 is not None
-    assert len(p3["code_predictor_codes"]) == _Q * 16
+    assert len(p3["codes"]["audio"]) == _Q * 16
 
 
 def test_ic_load_change_mid_request():
@@ -206,6 +207,7 @@ def test_ic_load_change_mid_request():
     assert _call(tm, "r", n_frames=27) is None
     p3 = _call(tm, "r", n_frames=49)
     assert p3 is not None
+    assert p3["meta"]["left_context_size"] == 24
 
     # A *new* request under high load gets IC=16 (not IC=2).
     # Frame 2 would emit under IC=2 but must hold under IC=16.
@@ -253,14 +255,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     assert payload is not None
-    assert payload["left_context_size"] == 2
-    assert len(payload["code_predictor_codes"]) == _Q * 12
+    assert payload["meta"]["left_context_size"] == 2
+    assert len(payload["codes"]["audio"]) == _Q * 12
 
 
 def test_ref_code_context_applies_to_all_streaming_chunks():
@@ -274,15 +276,15 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     assert payload is not None
     # ref_code (2 frames) prepended as left context on second chunk too
-    assert payload["left_context_size"] == 10 + 2
-    assert len(payload["code_predictor_codes"]) == _Q * (20 + 2)
+    assert payload["meta"]["left_context_size"] == 10 + 2
+    assert len(payload["codes"]["audio"]) == _Q * (20 + 2)
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
@@ -292,7 +294,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
 
     first_payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]]), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
@@ -302,22 +304,22 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     for _ in range(8):
         talker2code2wav_async_chunk(
             transfer_manager=tm,
-            pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+            pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
             request=_req(rid, finished=False, initial_codec_chunk_frames=10),
             is_finished=False,
         )
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     assert payload is not None
     # ref_code (2 frames) is kept (not popped) for subsequent chunks
-    assert payload["left_context_size"] == 2
-    assert len(payload["code_predictor_codes"]) == _Q * 12
+    assert payload["meta"]["left_context_size"] == 2
+    assert len(payload["codes"]["audio"]) == _Q * 12
     assert rid in tm.request_payload
 
 
@@ -332,7 +334,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         dtype=torch.long,
     )
     output = SimpleNamespace(
-        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(3)),
     )
     stage = SimpleNamespace(
@@ -343,7 +345,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
 
     assert len(prompts) == 1
     prompt = prompts[0]
-    assert prompt["additional_information"] == {"left_context_size": [2]}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": [2]}}
     assert prompt["prompt_token_ids"] == [
         9,
         8,
@@ -377,7 +379,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         dtype=torch.long,
     )
     output = SimpleNamespace(
-        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(4)),
     )
     stage = SimpleNamespace(
@@ -390,4 +392,4 @@ def test_non_async_processor_filters_out_of_range_codec_values():
     prompt = prompts[0]
     # Only ref_code (1 frame) + 2 valid frames = 3 frames * 4 quantizers = 12 codes
     assert len(prompt["prompt_token_ids"]) == 4 * 3
-    assert prompt["additional_information"] == {"left_context_size": [1]}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": [1]}}

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -126,8 +126,8 @@ def test_flush_tail_when_finished():
     )
 
     assert payload is not None
-    assert payload["finished"].item() is True
-    codes = payload["code_predictor_codes"]
+    assert payload["meta"]["finished"].item() is True
+    codes = payload["codes"]["audio"]
     # Format: [ctx_frames, context_length, ...flat_codes]
     assert len(codes) >= 2  # At least ctx_frames + context_length header
     ctx_frames = codes[0]
@@ -149,10 +149,8 @@ def test_eof_marker_when_finished_with_no_frames():
         request=request,
     )
 
-    assert payload == {
-        "code_predictor_codes": [],
-        "finished": torch.tensor(True, dtype=torch.bool),
-    }
+    assert payload["codes"] == {"audio": []}
+    assert payload["meta"]["finished"].item() is True
 
 
 def test_normal_chunk_emission():
@@ -176,7 +174,7 @@ def test_normal_chunk_emission():
 
     # A chunk should be emitted
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 20  # 25 - 5(chunk_size_at_begin)
@@ -203,7 +201,7 @@ def test_small_initial_chunks():
         )
 
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 0
@@ -252,7 +250,7 @@ def test_context_handling_format():
         )
 
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     # First two elements are ctx_frames and context_length
     ctx_frames = codes[0]
     context_length = codes[1]

tests/test_data_entry_keys.py

@@ -0,0 +1,252 @@
+"""Tests for data_entry_keys: TypedDict payload structure, flatten/unflatten, serialize/deserialize."""
+
+import torch
+
+from vllm_omni.data_entry_keys import (
+    OmniPayload,
+    deserialize_payload,
+    flatten_payload,
+    serialize_payload,
+    unflatten_payload,
+)
+from vllm_omni.engine import AdditionalInformationPayload
+
+
+class TestOmniPayload:
+    def test_nested_payload_structure(self):
+        """Verify OmniPayload can be constructed with nested dicts."""
+        payload: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0])},
+            "embed": {"prefill": torch.tensor([2.0])},
+            "codes": {"audio": torch.tensor([3.0])},
+            "ids": {"all": [1, 2, 3]},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+        }
+        assert torch.equal(payload["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(payload["embed"]["prefill"], torch.tensor([2.0]))
+        assert torch.equal(payload["codes"]["audio"], torch.tensor([3.0]))
+        assert payload["ids"]["all"] == [1, 2, 3]
+        assert payload["meta"]["finished"].item() is True
+
+    def test_partial_payload(self):
+        """OmniPayload fields are all optional (total=False)."""
+        payload: OmniPayload = {"meta": {"finished": torch.tensor(False, dtype=torch.bool)}}
+        assert payload["meta"]["finished"].item() is False
+
+    def test_empty_payload(self):
+        payload: OmniPayload = {}
+        assert len(payload) == 0
+
+
+class TestFlattenPayload:
+    def test_basic_nested_to_dotted(self):
+        nested = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["codes.audio"], torch.tensor([1.0]))
+        assert flat["meta.finished"].item() is True
+        assert flat["meta.left_context_size"] == 5
+        assert "codes" not in flat
+        assert "meta" not in flat
+
+    def test_top_level_keys_preserved(self):
+        nested = {
+            "latent": torch.tensor([9.0]),
+            "generated_len": 42,
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["latent"], torch.tensor([9.0]))
+        assert flat["generated_len"] == 42
+
+    def test_hidden_states_layers_expanded(self):
+        nested = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {
+                    0: torch.tensor([2.0]),
+                    24: torch.tensor([3.0]),
+                },
+            },
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["hidden_states.output"], torch.tensor([1.0]))
+        assert torch.equal(flat["hidden_states.layer_0"], torch.tensor([2.0]))
+        assert torch.equal(flat["hidden_states.layer_24"], torch.tensor([3.0]))
+        assert "hidden_states.layers" not in flat
+
+    def test_empty_payload(self):
+        assert flatten_payload({}) == {}
+
+    def test_mixed_nested_and_top_level(self):
+        nested: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "latent": torch.tensor([2.0]),
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        }
+        flat = flatten_payload(nested)
+        assert set(flat.keys()) == {"codes.audio", "latent", "meta.finished"}
+
+
+class TestUnflattenPayload:
+    def test_basic_dotted_to_nested(self):
+        flat = {
+            "codes.audio": torch.tensor([1.0]),
+            "meta.finished": torch.tensor(True, dtype=torch.bool),
+            "meta.left_context_size": 5,
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["codes"]["audio"], torch.tensor([1.0]))
+        assert nested["meta"]["finished"].item() is True
+        assert nested["meta"]["left_context_size"] == 5
+
+    def test_top_level_keys_preserved(self):
+        flat = {"latent": torch.tensor([9.0]), "generated_len": 42}
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["latent"], torch.tensor([9.0]))
+        assert nested["generated_len"] == 42
+
+    def test_hidden_states_layers_collected(self):
+        flat = {
+            "hidden_states.output": torch.tensor([1.0]),
+            "hidden_states.layer_0": torch.tensor([2.0]),
+            "hidden_states.layer_24": torch.tensor([3.0]),
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_empty_payload(self):
+        assert unflatten_payload({}) == {}
+
+
+class TestFlattenUnflattenRoundTrip:
+    def test_round_trip_simple(self):
+        original: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0, 2.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 10},
+            "ids": {"prompt": [1, 2, 3]},
+            "latent": torch.tensor([5.0]),
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 10
+        assert restored["ids"]["prompt"] == [1, 2, 3]
+        assert torch.equal(restored["latent"], original["latent"])
+
+    def test_round_trip_with_layers(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_round_trip_all_categories(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0]), "last": torch.tensor([2.0])},
+            "embed": {"prefill": torch.tensor([3.0]), "tts_bos": torch.tensor([4.0])},
+            "codes": {"audio": torch.tensor([5.0]), "ref": torch.tensor([6.0])},
+            "ids": {"all": [1, 2], "prompt": [3, 4]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 8},
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["last"], torch.tensor([2.0]))
+        assert torch.equal(restored["embed"]["prefill"], torch.tensor([3.0]))
+        assert torch.equal(restored["embed"]["tts_bos"], torch.tensor([4.0]))
+        assert torch.equal(restored["codes"]["audio"], torch.tensor([5.0]))
+        assert torch.equal(restored["codes"]["ref"], torch.tensor([6.0]))
+        assert restored["ids"]["all"] == [1, 2]
+        assert restored["ids"]["prompt"] == [3, 4]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 8
+
+
+class TestSerializeDeserializePayload:
+    def test_tensor_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
+        }
+        wire = serialize_payload(original)
+        assert isinstance(wire, AdditionalInformationPayload)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+
+    def test_list_round_trip(self):
+        original: OmniPayload = {
+            "ids": {"prompt": [10, 20, 30]},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["ids"]["prompt"] == [10, 20, 30]
+
+    def test_finished_tensor_round_trip(self):
+        original: OmniPayload = {
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert isinstance(restored["meta"]["finished"], torch.Tensor)
+        assert restored["meta"]["finished"].dtype == torch.bool
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 5
+
+    def test_mixed_types_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0, 2.0])},
+            "ids": {"all": [1, 2, 3]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4},
+            "codes": {"audio": torch.tensor([3.0])},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+        assert restored["ids"]["all"] == [1, 2, 3]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 4
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+
+    def test_hidden_states_layers_round_trip(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_tensor_dtype_preserved(self):
+        # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
+        for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
+            original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
+            wire = serialize_payload(original)
+            restored = deserialize_payload(wire)
+            assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
+
+    def test_tensor_shape_preserved(self):
+        t = torch.randn(3, 4, 5)
+        original: OmniPayload = {"hidden_states": {"output": t}}
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["hidden_states"]["output"].shape == (3, 4, 5)
+        assert torch.allclose(restored["hidden_states"]["output"], t)
+
+    def test_empty_payload_returns_none(self):
+        assert serialize_payload({}) is None
+
+    def test_none_values_skipped(self):
+        original: OmniPayload = {"meta": {"finished": None}}
+        wire = serialize_payload(original)
+        assert wire is None

tests/worker/test_omni_connector_mixin.py

@@ -891,9 +891,8 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
     def test_rank0_only_polls_connector_for_tp_async_chunk(self):
         host = self._make_host(rank=0)
         payload = {
-            "code_predictor_codes": [10, 11],
-            "left_context_size": 0,
-            "finished": torch.tensor(False),
+            "codes": {"audio": [10, 11]},
+            "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
         }
         host._omni_connector.get.return_value = (payload, 123)
         tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
@@ -1165,9 +1164,8 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
 
         host._omni_connector.get.return_value = (
             {
-                "code_predictor_codes": [20, 21, 22],
-                "left_context_size": 0,
-                "finished": torch.tensor(False),
+                "codes": {"audio": [20, 21, 22]},
+                "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
             },
             1,
         )

tests/worker/test_omni_gpu_model_runner.py

@@ -119,7 +119,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
 
     runner.talker_mtp = DummyTalkerMTP()
-    runner.model = SimpleNamespace(talker_mtp_output_key="code_predictor_codes")
+    runner.model = SimpleNamespace(talker_mtp_output_key=("codes", "audio"))
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace())
 
     # Provide a minimal implementation that returns the expected 4-tuple.
@@ -188,8 +188,8 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
     # Validate per-request additional_information_cpu was updated
     info_r1 = runner.requests["r1"].additional_information_cpu
     info_r2 = runner.requests["r2"].additional_information_cpu
-    assert int(info_r1["code_predictor_codes"][0, 0]) == 0
-    assert int(info_r2["code_predictor_codes"][0, 0]) == 1
+    assert int(info_r1["codes"]["audio"][0, 0]) == 0
+    assert int(info_r2["codes"]["audio"][0, 0]) == 1
 
 
 def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -314,7 +314,7 @@ class OmniSchedulingCoordinator:
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
-                        "left_context_size": metadata["left_context_size"],
+                        "meta": {"left_context_size": metadata["left_context_size"]},
                     }
                 request._omni_initial_model_buffer = runtime_seed
                 if new_ids:

vllm_omni/data_entry_keys.py

@@ -0,0 +1,277 @@
+"""Structured payload types for inter-stage communication.
+
+Adding a new model?
+~~~~~~~~~~~~~~~~~~~
+Every key you put into the inter-stage payload (``additional_information``,
+``multimodal_output``, ``pooling_output``) **must** use the nested
+``OmniPayload`` TypedDict structure.  For each category, every known
+qualifier is an explicit field so misspellings are caught statically.
+
+Categories
+    hidden_states  – intermediate / output hidden-state tensors
+    embed          – embedding tensors (prefill, decode, special tokens)
+    ids            – token-ID sequences
+    codes          – codec / audio code tensors
+    meta           – scalar metadata, control flags, shapes
+
+This module provides:
+- Structured ``TypedDict`` types for static type checking (``OmniPayload``)
+- ``serialize_payload`` / ``deserialize_payload`` for transport across
+  process boundaries via ``AdditionalInformationPayload``
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from vllm_omni.engine import AdditionalInformationEntry, AdditionalInformationPayload
+
+# ── Structured payload types ──
+# These are TypedDicts (plain dicts at runtime, zero overhead) that give
+# static type checking and IDE autocomplete for inter-stage payloads.
+# Every field is optional (total=False) because each stage only populates
+# the subset it needs.
+
+
+class HiddenStates(TypedDict, total=False):
+    output: torch.Tensor
+    trailing_text: torch.Tensor
+    last: torch.Tensor
+    layers: dict[int, torch.Tensor]
+
+
+class Embeddings(TypedDict, total=False):
+    prefill: torch.Tensor
+    decode: torch.Tensor
+    cached_decode: torch.Tensor
+    tts_bos: torch.Tensor
+    tts_eos: torch.Tensor
+    tts_pad: torch.Tensor
+    tts_pad_projected: torch.Tensor
+    voice: torch.Tensor
+    speech_feat: torch.Tensor
+    thinker_reply: torch.Tensor
+
+
+class Codes(TypedDict, total=False):
+    audio: torch.Tensor
+    ref: torch.Tensor
+
+
+class Ids(TypedDict, total=False):
+    all: list[int]
+    prompt: list[int]
+    output: list[int]
+    speech_token: list[int]
+    prior_image: list[int]
+
+
+class OmniPayloadMeta(TypedDict, total=False):
+    finished: torch.Tensor
+    left_context_size: int
+    override_keys: list[tuple[str, str]]
+    num_processed_tokens: int
+    next_stage_prompt_len: int
+    ar_width: int
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    gen_token_mask: torch.Tensor
+    omni_task: list[str]
+    height: int
+    width: int
+    decode_flag: bool
+    codec_streaming: bool
+    ref_code_len: int
+    talker_prefill_offset: int
+
+
+class OmniPayload(TypedDict, total=False):
+    hidden_states: HiddenStates
+    embed: Embeddings
+    ids: Ids
+    codes: Codes
+    meta: OmniPayloadMeta
+    latent: torch.Tensor
+    generated_len: int
+    model_outputs: list[torch.Tensor]
+    mtp_inputs: tuple[torch.Tensor, torch.Tensor]
+    speaker: Any
+    language: Any
+    request_id: str
+
+
+# ── Keys whose values are nested dicts (TypedDict sub-categories) ──
+_NESTED_KEYS = frozenset({"hidden_states", "embed", "ids", "codes", "meta"})
+
+# Sub-TypedDict for each nested category, used by runtime validation.
+_NESTED_SCHEMAS: dict[str, type] = {
+    "hidden_states": HiddenStates,
+    "embed": Embeddings,
+    "ids": Ids,
+    "codes": Codes,
+    "meta": OmniPayloadMeta,
+}
+
+_ROOT_KEYS: frozenset[str] = frozenset(OmniPayload.__annotations__.keys())
+
+
+def assert_payload(payload: dict[str, Any], *, context: str = "payload") -> None:
+    """Validate ``payload`` matches the ``OmniPayload`` nested schema.
+
+    TypedDict is a static-only contract in Python; this helper closes the
+    loop at runtime by rejecting:
+      * non-dict payloads
+      * top-level keys not declared on ``OmniPayload``
+      * nested-category values that aren't dicts
+      * sub-keys not declared on the matching nested TypedDict
+
+    Call at producer/consumer boundaries when a schema violation should
+    crash the pipeline instead of silently degrading audio quality.
+    """
+    assert isinstance(payload, dict), f"{context}: expected dict, got {type(payload).__name__}"
+    extra_top = set(payload) - _ROOT_KEYS
+    assert not extra_top, f"{context}: unknown top-level keys {sorted(extra_top)!r}"
+    for nested_key, schema in _NESTED_SCHEMAS.items():
+        if nested_key not in payload:
+            continue
+        sub = payload[nested_key]
+        assert isinstance(sub, dict), f"{context}: payload[{nested_key!r}] must be dict, got {type(sub).__name__}"
+        known_sub = frozenset(schema.__annotations__.keys())
+        extra_sub = set(sub) - known_sub
+        assert not extra_sub, f"{context}: payload[{nested_key!r}] unknown sub-keys {sorted(extra_sub)!r}"
+
+
+def flatten_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    """Flatten a nested ``OmniPayload`` to dotted keys.
+
+    Nested sub-dicts under ``_NESTED_KEYS`` are expanded:
+    ``{"codes": {"audio": tensor}}`` → ``{"codes.audio": tensor}``.
+    ``hidden_states["layers"]`` is expanded to ``hidden_states.layer_N``.
+    Top-level values are kept as-is.
+    """
+    if not payload:
+        return {}
+    flat: dict[str, Any] = {}
+    for key, value in payload.items():
+        if key in _NESTED_KEYS and isinstance(value, dict):
+            for qual, val in value.items():
+                if qual == "layers" and key == "hidden_states" and isinstance(val, dict):
+                    for layer_idx, tensor in val.items():
+                        flat[f"hidden_states.layer_{layer_idx}"] = tensor
+                else:
+                    flat[f"{key}.{qual}"] = val
+        else:
+            flat[key] = value
+    return flat
+
+
+def unflatten_payload(flat: dict[str, Any]) -> dict[str, Any]:
+    """Unflatten dotted keys back to nested dicts.
+
+    Reverse of :func:`flatten_payload`.
+    ``hidden_states.layer_N`` keys are collected into ``hidden_states.layers``.
+    """
+    result: dict[str, Any] = {}
+    for key, value in flat.items():
+        if "." in key:
+            type_key, qualifier = key.split(".", 1)
+            sub = result.setdefault(type_key, {})
+            if type_key == "hidden_states" and qualifier.startswith("layer_"):
+                layers = sub.setdefault("layers", {})
+                layer_idx = int(qualifier[len("layer_") :])
+                layers[layer_idx] = value
+            else:
+                sub[qualifier] = value
+        else:
+            result[key] = value
+    return result
+
+
+# ── dtype helpers ──
+_DTYPE_TO_NAME: dict[torch.dtype, str] = {
+    torch.float32: "float32",
+    torch.float16: "float16",
+    torch.bfloat16: "bfloat16",
+    torch.float64: "float64",
+    torch.int64: "int64",
+    torch.int32: "int32",
+    torch.int16: "int16",
+    torch.int8: "int8",
+    torch.uint8: "uint8",
+    torch.bool: "bool",
+}
+
+
+def _dtype_to_name(dtype: torch.dtype) -> str:
+    return _DTYPE_TO_NAME.get(dtype, str(dtype).replace("torch.", ""))
+
+
+def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry:
+    from vllm_omni.engine import AdditionalInformationEntry
+
+    t_cpu = t.detach().to("cpu").contiguous()
+    return AdditionalInformationEntry(
+        tensor_data=t_cpu.numpy().tobytes(),
+        tensor_shape=list(t_cpu.shape),
+        tensor_dtype=_dtype_to_name(t_cpu.dtype),
+    )
+
+
+def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor:
+    dt = np.dtype(entry.tensor_dtype or "float32")
+    arr = np.frombuffer(entry.tensor_data, dtype=dt)  # type: ignore[arg-type]
+    arr = arr.reshape(entry.tensor_shape)
+    return torch.from_numpy(arr.copy())
+
+
+def serialize_payload(
+    payload: OmniPayload,
+) -> AdditionalInformationPayload | None:
+    """Serialize an ``OmniPayload`` for EngineCore transport.
+
+    Uses :func:`flatten_payload` to produce dotted keys, then converts
+    each value to an ``AdditionalInformationEntry``.
+    """
+    from vllm_omni.engine import (
+        AdditionalInformationEntry,
+        AdditionalInformationPayload,
+    )
+
+    flat = flatten_payload(payload)
+    entries: dict[str, AdditionalInformationEntry] = {}
+
+    for key, value in flat.items():
+        if isinstance(value, torch.Tensor):
+            entries[key] = _serialize_tensor(value)
+        elif isinstance(value, list):
+            entries[key] = AdditionalInformationEntry(list_data=value)
+        elif value is not None:
+            entries[key] = AdditionalInformationEntry(scalar_data=value)
+
+    return AdditionalInformationPayload(entries=entries) if entries else None
+
+
+def deserialize_payload(
+    wire: AdditionalInformationPayload,
+) -> OmniPayload:
+    """Deserialize an ``AdditionalInformationPayload`` back to ``OmniPayload``.
+
+    Decodes entries to tensors/lists, then uses :func:`unflatten_payload`
+    to reconstruct the nested structure.
+    """
+    flat: dict[str, Any] = {}
+
+    for key, entry in wire.entries.items():
+        if entry.tensor_data is not None:
+            flat[key] = _deserialize_tensor(entry)
+        elif entry.list_data is not None:
+            flat[key] = entry.list_data
+        elif entry.scalar_data is not None:
+            flat[key] = entry.scalar_data
+
+    return unflatten_payload(flat)  # type: ignore[return-value]

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -8,6 +8,8 @@ from typing import Any
 import torch
 from vllm.v1.request import Request, RequestStatus
 
+from vllm_omni.data_entry_keys import unflatten_payload
+
 from ..factory import OmniConnectorFactory
 from ..utils.config import ConnectorSpec
 from ..utils.logging import get_connector_logger
@@ -150,30 +152,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             # Update connector state
             self.get_req_chunk[req_id] += 1
 
+            meta = payload_data.get("meta", {})
             if self.model_mode == "ar":
-                self._update_request_payload(external_req_id, payload_data)
-                request.additional_information = payload_data
-                if payload_data.get("finished"):
+                merged_payload = self._update_request_payload(external_req_id, payload_data)
+                request.additional_information = merged_payload
+                if meta.get("finished"):
                     self.finished_requests.add(req_id)
             else:
-                if payload_data.get("finished"):
+                if meta.get("finished"):
                     self.finished_requests.add(req_id)
 
-                new_ids = payload_data.get("code_predictor_codes", [])
+                new_ids = payload_data.get("codes", {}).get("audio", [])
                 request.prompt_token_ids = new_ids
-                # Preserve previously attached request metadata (e.g. prompt
-                # conditioning tensors) and update only per-chunk fields.
                 prev_info = getattr(request, "additional_information", None)
                 info = dict(prev_info) if isinstance(prev_info, dict) else {}
                 for key, value in payload_data.items():
-                    if key in {"code_predictor_codes", "finished"}:
+                    if key == "codes":
+                        continue
+                    if isinstance(value, dict):
+                        existing_sub = info.get(key)
+                        merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
+                        for sk, sv in value.items():
+                            if key == "meta" and sk == "finished":
+                                continue
+                            merged_sub[sk] = sv
+                        info[key] = merged_sub
                         continue
                     info[key] = value
                 request.additional_information = info
                 request.num_computed_tokens = 0
 
                 # Empty chunk with more data expected: keep polling.
-                if not new_ids and not payload_data.get("finished"):
+                if not new_ids and not meta.get("finished"):
                     return True
 
             # Mark as finished for consumption
@@ -184,33 +194,36 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         return False
 
     def _update_request_payload(self, req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]:
-        """Update the payload data for a request in the connector.
-
-        Args:
-            connector: OmniConnectorBase instance
-            req_id: Request ID to update
-            payload_data: New payload data to store
-        """
+        """Update the stored payload for *req_id* with the latest chunk."""
         if req_id not in self.request_payload:
             self.request_payload[req_id] = payload_data
             return payload_data
-        origin_payload = self.request_payload[req_id]
-        override_keys = payload_data.pop("override_keys", [])
-        for key, value in payload_data.items():
-            if key == "finished":
+        origin = self.request_payload[req_id]
+        raw_ok = payload_data.get("meta", {}).pop("override_keys", [])
+        override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok}
+
+        for type_key, new_val in payload_data.items():
+            if not isinstance(new_val, dict):
                 continue
-            elif key in override_keys:
-                payload_data[key] = value
-            elif isinstance(value, torch.Tensor) and key in origin_payload:
-                payload_data[key] = torch.cat([origin_payload[key], value], dim=0)
-            elif isinstance(value, list) and key in origin_payload:
-                payload_data[key] = origin_payload[key] + value
+            origin_sub = origin.get(type_key)
+            if not isinstance(origin_sub, dict):
+                continue
+            for qual, value in new_val.items():
+                if type_key == "meta" and qual == "finished":
+                    continue
+                if (type_key, qual) in override_keys:
+                    continue
+                if isinstance(value, torch.Tensor) and qual in origin_sub:
+                    new_val[qual] = torch.cat([origin_sub[qual], value], dim=0)
+                elif isinstance(value, list) and qual in origin_sub:
+                    new_val[qual] = origin_sub[qual] + value
 
         self.request_payload[req_id] = payload_data
         return payload_data
 
     def _send_single_request(self, task: dict):
-        pooling_output = task["pooling_output"]
+        raw_po = task["pooling_output"]
+        pooling_output = unflatten_payload(raw_po) if isinstance(raw_po, dict) else raw_po
         request = task["request"]
         is_finished = task["is_finished"]
         stage_id = self.connector.stage_id
@@ -245,7 +258,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if success:
             self.put_req_chunk[external_req_id] += 1
             logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}")
-            finished_flag = payload_data.get("finished")
+            finished_flag = payload_data.get("meta", {}).get("finished", payload_data.get("finished"))
             is_payload_finished = False
             if isinstance(finished_flag, torch.Tensor):
                 is_payload_finished = finished_flag.numel() == 1 and bool(finished_flag.item())

vllm_omni/engine/output_processor.py

@@ -16,6 +16,7 @@ from vllm.v1.engine.output_processor import (
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats
 
+from vllm_omni.data_entry_keys import unflatten_payload
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -56,23 +57,21 @@ class OmniRequestState(RequestState):
                 return x
 
             if isinstance(payload, dict):
+                # Keep payload flat (dotted keys like "hidden_states.layer_0")
+                # during accumulation so that all values are tensors/scalars and
+                # the merge logic below works correctly.  Unflatten happens
+                # later in _consolidate_multimodal_tensors after concatenation.
+
                 incoming: dict[str, Any] = {}
                 target_key = self.mm_type or "hidden"
 
                 for k, v in payload.items():
-                    # Normalize producer keys to the modality name.
-                    # AR runners produce {"hidden": ...} and generation
-                    # runners produce {"model_outputs": ...}; remap both
-                    # to the semantic modality key (e.g. "audio", "latent").
                     if k == "model_outputs":
                         k = target_key
                     elif k == "hidden" and target_key != "hidden":
                         k = target_key
 
-                    if isinstance(v, dict):
-                        incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()}
-                    else:
-                        incoming[k] = _to_cpu(v)
+                    incoming[k] = _to_cpu(v)
             else:
                 key = self.mm_type or "hidden"
                 incoming = {key: _to_cpu(payload)}
@@ -80,20 +79,16 @@ class OmniRequestState(RequestState):
             if self.mm_accumulated is None:
                 self.mm_accumulated = incoming
             else:
-                # Merge keys; accumulate tensors in lists for deferred concatenation
                 for k, v in incoming.items():
                     if k not in self.mm_accumulated:
                         self.mm_accumulated[k] = v
                     else:
                         existing = self.mm_accumulated[k]
                         if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
-                            # Use list accumulation to avoid O(n²) repeated concatenation
                             self.mm_accumulated[k] = [existing, v]
                         elif isinstance(v, torch.Tensor) and isinstance(existing, list):
-                            # Append to existing list
                             existing.append(v)
                         elif isinstance(v, dict) and isinstance(existing, dict):
-                            # Merge nested dicts with list accumulation for tensors
                             for sk, sv in v.items():
                                 if sk not in existing:
                                     existing[sk] = sv
@@ -140,6 +135,13 @@ class OmniRequestState(RequestState):
         except Exception:
             logger.exception("Error consolidating multimodal tensors")
 
+        # Restore nested structure from flat dotted keys now that all tensor
+        # lists have been concatenated into single tensors.
+        try:
+            self.mm_accumulated = unflatten_payload(self.mm_accumulated)
+        except Exception:
+            logger.exception("Error unflattening consolidated multimodal tensors")
+
     # Override: do not route to pooling-only path; always create completion
     # outputs, and attach pooling_result into the CompletionOutput.
     def make_request_output(

vllm_omni/engine/serialization.py

@@ -4,110 +4,39 @@ from __future__ import annotations
 
 from typing import Any
 
-import numpy as np
-import torch
 from vllm.logger import init_logger
 
-from vllm_omni.engine import (
-    AdditionalInformationEntry,
-    AdditionalInformationPayload,
-)
+from vllm_omni.data_entry_keys import OmniPayload, deserialize_payload, serialize_payload
+from vllm_omni.engine import AdditionalInformationPayload
 
 logger = init_logger(__name__)
 
 
-def dtype_to_name(dtype: torch.dtype) -> str:
-    """Convert torch dtype to a stable string name for serialization."""
-    mapping = {
-        torch.float32: "float32",
-        torch.float: "float32",
-        torch.float16: "float16",
-        torch.half: "float16",
-        torch.bfloat16: "bfloat16",
-        torch.float64: "float64",
-        torch.double: "float64",
-        torch.int64: "int64",
-        torch.long: "int64",
-        torch.int32: "int32",
-        torch.int: "int32",
-        torch.int16: "int16",
-        torch.short: "int16",
-        torch.int8: "int8",
-        torch.uint8: "uint8",
-        torch.bool: "bool",
-    }
-    return mapping.get(dtype, str(dtype).replace("torch.", ""))
-
-
 def serialize_additional_information(
     raw_info: dict[str, Any] | AdditionalInformationPayload | None,
     *,
     log_prefix: str | None = None,
 ) -> AdditionalInformationPayload | None:
-    """Serialize omni request metadata for EngineCore transport."""
+    """Serialize omni request metadata for EngineCore transport.
+
+    Delegates to ``serialize_payload`` which understands the nested
+    ``OmniPayload`` TypedDict structure.
+    """
     if raw_info is None:
         return None
     if isinstance(raw_info, AdditionalInformationPayload):
         return raw_info
 
-    entries: dict[str, AdditionalInformationEntry] = {}
-    for key, value in raw_info.items():
-        if isinstance(value, torch.Tensor):
-            value_cpu = value.detach().to("cpu").contiguous()
-            entries[key] = AdditionalInformationEntry(
-                tensor_data=value_cpu.numpy().tobytes(),
-                tensor_shape=list(value_cpu.shape),
-                tensor_dtype=dtype_to_name(value_cpu.dtype),
-            )
-            continue
-
-        if isinstance(value, list):
-            entries[key] = AdditionalInformationEntry(list_data=value)
-            continue
-
-        entries[key] = AdditionalInformationEntry(scalar_data=value)
-
-    return AdditionalInformationPayload(entries=entries) if entries else None
+    payload: OmniPayload = raw_info  # type: ignore[assignment]
+    return serialize_payload(payload)
 
 
 def deserialize_additional_information(
-    payload: dict | AdditionalInformationPayload | object | None,
+    payload: dict | AdditionalInformationPayload | None,
 ) -> dict:
-    """Deserialize an *additional_information* payload into a plain dict.
-
-    Accepts:
-    - ``dict`` – returned as-is.
-    - ``AdditionalInformationPayload`` (or duck-typed with
-      ``.entries``) – decoded entry-by-entry.
-    - ``None`` – returns ``{}``.
-    """
-
+    """Deserialize an *additional_information* payload into a plain dict."""
     if payload is None:
         return {}
-
     if isinstance(payload, dict):
         return payload
-
-    try:
-        entries = getattr(payload, "entries", None)
-        if not isinstance(entries, dict):
-            logger.exception("Failed to decode additional_information payload, entries field not a dict")
-            return {}
-        info: dict[str, object] = {}
-        for k, entry in entries.items():
-            if getattr(entry, "tensor_data", None) is not None:
-                dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
-                arr = np.frombuffer(entry.tensor_data, dtype=dt)
-                arr = arr.reshape(getattr(entry, "tensor_shape", ()))
-                info[k] = torch.from_numpy(arr.copy())
-            elif getattr(entry, "list_data", None) is not None:
-                info[k] = entry.list_data
-            elif getattr(entry, "scalar_data", None) is not None:
-                info[k] = entry.scalar_data
-            else:
-                info[k] = None
-        return info
-    except Exception:
-        logger.exception("Failed to decode additional_information payload")
-
-    return {}
+    return deserialize_payload(payload)  # type: ignore[return-value]

vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py

@@ -230,8 +230,9 @@ class FishSpeechDACDecoder(nn.Module):
             for i, info in enumerate(runtime_additional_information):
                 if i >= len(left_context_size):
                     break
-                if "left_context_size" in info:
-                    left_context_size[i] = info["left_context_size"]
+                meta = info.get("meta", {}) if isinstance(info, dict) else {}
+                if "left_context_size" in meta:
+                    left_context_size[i] = meta["left_context_size"]
 
         for i, req_ids in enumerate(request_ids_list):
             if req_ids.numel() < 1:

vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py

@@ -193,9 +193,9 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         self.has_preprocess = True
         self.has_postprocess = True
         self.mtp_hidden_size = int(self.text_config.hidden_size)
-        self.talker_mtp_output_key = "audio_codes"
+        self.talker_mtp_output_key = ("codes", "audio")
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")}
         self.talker_mtp_graph_safe = True
-        self.gpu_resident_buffer_keys: set[str] = {"last_slow_ar_hidden"}
 
         # Qwen3 transformer backbone.
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
@@ -330,7 +330,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         for info in info_dicts:
             if not isinstance(info, dict):
                 continue
-            ac = info.get("audio_codes")
+            ac = info.get("codes", {}).get("audio")
             if isinstance(ac, torch.Tensor):
                 audio_codes_list.append(ac)
 
@@ -365,7 +365,8 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
 
         if span_len > 1:
             # --- Prefill ---
-            prompt_embeds_buf = info_dict.get("slow_ar_prompt_embeds")
+            embed = info_dict.get("embed", {})
+            prompt_embeds_buf = embed.get("prefill")
             is_first_prefill = not isinstance(prompt_embeds_buf, torch.Tensor) or prompt_embeds_buf.ndim != 2
             dev = input_ids.device
 
@@ -381,8 +382,8 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
                 next_offset = min(span_len, total_prompt_len)
 
                 info_update: dict[str, Any] = {
-                    "slow_ar_prompt_embeds": prompt_embeds_buf if next_offset < total_prompt_len else None,
-                    "prefill_offset": next_offset,
+                    "embed": {"prefill": prompt_embeds_buf if next_offset < total_prompt_len else None},
+                    "meta": {"prefill_offset": next_offset},
                 }
 
                 take = prompt_embeds_buf[:span_len]
@@ -399,7 +400,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
                     device=dev,
                     dtype=torch.long,
                 )
-                info_update["audio_codes"] = zeros
+                info_update["codes"] = {"audio": zeros}
 
                 input_ids_out = input_ids.clone()
                 input_ids_out[:] = self._audio_pad_token_id
@@ -407,7 +408,8 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
 
             else:
                 # Subsequent prefill chunk.
-                offset = int(info_dict.get("prefill_offset", 0) or 0)
+                meta = info_dict.get("meta", {})
+                offset = int(meta.get("prefill_offset", 0) or 0)
                 total_prompt_len = int(prompt_embeds_buf.shape[0])
                 s = max(0, min(offset, total_prompt_len))
                 e = max(0, min(offset + span_len, total_prompt_len))
@@ -426,20 +428,21 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
                     input_ids.clone().fill_(self._audio_pad_token_id),
                     prompt_embeds,
                     {
-                        "slow_ar_prompt_embeds": prompt_embeds_buf if next_offset < total_prompt_len else None,
-                        "prefill_offset": next_offset,
-                        "audio_codes": zeros,
+                        "embed": {"prefill": prompt_embeds_buf if next_offset < total_prompt_len else None},
+                        "meta": {"prefill_offset": next_offset},
+                        "codes": {"audio": zeros},
                     },
                 )
 
         # --- Decode: span_len == 1 ---
         dev = input_ids.device
 
-        last_hidden = info_dict.get("last_slow_ar_hidden")
+        hs = info_dict.get("hidden_states", {})
+        last_hidden = hs.get("last")
         if not isinstance(last_hidden, torch.Tensor):
             # First decode step after prefill -- just embed the token directly.
             logger.warning(
-                "preprocess decode: last_slow_ar_hidden not found (keys=%s), "
+                "preprocess decode: hidden_states.last not found (keys=%s), "
                 "returning plain embed (mtp_inputs will NOT be set)",
                 list(info_dict.keys()),
             )
@@ -468,8 +471,8 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
             logger.debug("postprocess: empty hidden_states")
             return {}
         last = hidden_states[-1, :].detach().contiguous()
-        logger.debug("postprocess: saved last_slow_ar_hidden shape=%s", tuple(last.shape))
-        return {"last_slow_ar_hidden": last}
+        logger.debug("postprocess: saved hidden_states.last shape=%s", tuple(last.shape))
+        return {"hidden_states": {"last": last}}
 
     # -------------------- prompt construction --------------------
 

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -2321,7 +2321,7 @@ class GlmImageModel(nn.Module):
                 upsampled_token_ids.append(tokens_upsampled.view(-1))
 
             prior_token_image_ids_info = {
-                "prior_token_image_ids": upsampled_token_ids,
+                "ids": {"prior_image": upsampled_token_ids},
                 "image_grid_thw": image_grid_thw.tolist(),
             }
 
@@ -2533,9 +2533,7 @@ class GlmImageForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP
         # image_grid_thw is NOT included because:
         # 1. vLLM's pooling_output expects dict[str, torch.Tensor], not mixed types
         # 2. ar2diffusion doesn't need it - the grid info is already encoded in tensor shape
-        prior_token_info = {
-            "prior_token_image_ids": upsampled_token_ids,
-        }
+        prior_token_info = {"ids": {"prior_image": upsampled_token_ids}}
 
         # Debug: log prior_token_info
         shapes = [t.shape for t in upsampled_token_ids]

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -608,16 +608,17 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration)
         num_reqs = int(logits.shape[0])
         for i in range(num_reqs):
             runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {}
-            omni_task = runtime_info.get("omni_task")
+            meta = runtime_info.get("meta", {})
+            omni_task = meta.get("omni_task")
             if not isinstance(omni_task, list) or not omni_task or omni_task[0] != "t2i":
                 # Text/understanding/chat: forbid sampling from the extra gen vocab.
                 logits[i, self.language_model.base_vocab_size :] = neg_inf
                 continue
 
-            ar_width = runtime_info["ar_width"][0]
-            eol_token_id = runtime_info["eol_token_id"][0]
-            visual_start = runtime_info["visual_token_start_id"][0]
-            visual_end = runtime_info["visual_token_end_id"][0]
+            ar_width = meta["ar_width"][0]
+            eol_token_id = meta["eol_token_id"][0]
+            visual_start = meta["visual_token_start_id"][0]
+            visual_end = meta["visual_token_end_id"][0]
             generated_len = runtime_info["generated_len"]
 
             row = logits[i]

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -799,7 +799,7 @@ class MiMoAudioForConditionalGeneration(
 
             return OmniOutput(
                 text_hidden_states=text_hidden_states.reshape(-1, text_hidden_states.shape[-1]),
-                multimodal_outputs={"code_predictor_codes": next_speech_tokens},
+                multimodal_outputs={"codes": {"audio": next_speech_tokens}},
             )
 
         if self.model_stage == "code2wav":

vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py

@@ -486,10 +486,10 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             return [ids]
 
         if runtime_additional_information and all(
-            isinstance(info.get("code_flat_numel"), int) and int(info["code_flat_numel"]) > 0
+            isinstance(info.get("meta", {}).get("code_flat_numel"), int) and int(info["meta"]["code_flat_numel"]) > 0
             for info in runtime_additional_information
         ):
-            sizes = [int(info["code_flat_numel"]) for info in runtime_additional_information]
+            sizes = [int(info["meta"]["code_flat_numel"]) for info in runtime_additional_information]
             if sum(sizes) == n:
                 parts: list[torch.Tensor] = []
                 offset = 0
@@ -517,11 +517,11 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         if not runtime_additional_information:
             return left_frames, chunk_frames
         for i in range(min(num_req, len(runtime_additional_information))):
-            info = runtime_additional_information[i]
-            if "left_context_size" in info:
-                left_frames[i] = int(info["left_context_size"])
-            if "codec_chunk_frames" in info:
-                chunk_frames[i] = int(info["codec_chunk_frames"])
+            meta = runtime_additional_information[i].get("meta", {})
+            if "left_context_size" in meta:
+                left_frames[i] = int(meta["left_context_size"])
+            if "codec_chunk_frames" in meta:
+                chunk_frames[i] = int(meta["codec_chunk_frames"])
         return left_frames, chunk_frames
 
     def chunked_decode_streaming(

vllm_omni/model_executor/models/output_templates.py

@@ -3,11 +3,13 @@ from typing import NamedTuple
 import torch
 from vllm.sequence import IntermediateTensors
 
+from vllm_omni.data_entry_keys import OmniPayload
+
 
 class OmniOutput(NamedTuple):
     """Output from the merged Omni model containing both text and audio."""
 
     text_hidden_states: torch.Tensor
-    multimodal_outputs: dict | None = None
+    multimodal_outputs: OmniPayload | None = None
     intermediate_tensors: IntermediateTensors | None = None
     next_token_id: torch.Tensor | None = None

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -32,6 +32,7 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.model_executor.models.output_templates import OmniOutput
@@ -666,6 +667,8 @@ class Qwen2_5OmniForConditionalGeneration(
         # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional),
         #   add them; otherwise, keep the original embedding.
 
+        payload: OmniPayload = info_dict
+
         # Ensure we have base embeddings when only ids are provided
         if input_embeds is None and input_ids is not None:
             input_embeds = self.talker.embed_input_ids(input_ids)
@@ -673,23 +676,27 @@ class Qwen2_5OmniForConditionalGeneration(
         span_len = input_ids.shape[0]
         if span_len > 1:
             # prefill
-            return self.thinker_to_talker_process(input_ids, input_embeds, **info_dict)
+            return self.thinker_to_talker_process(input_ids, input_embeds, payload)
         else:
             # decode
-            return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, **info_dict)
+            return self.thinker_to_talker_decode_one_step(input_ids, input_embeds, payload)
 
     def thinker_to_talker_process(
         self,
         input_ids: torch.Tensor,
         input_embeds: torch.Tensor,
-        **info_dict: object,
+        payload: OmniPayload,
     ):
+        embed = payload.get("embed", {})
+        hs = payload.get("hidden_states", {})
+        ids = payload.get("ids", {})
+
         update_dict = {}
 
-        prompt_embeds = info_dict.get("prompt_embeds")  # Tensor [P,H]
-        thinker_result = info_dict.get("thinker_result")  # Tensor [K,H]
-        prompt_token_ids = info_dict.get("prompt_token_ids")  # list[int]
-        thinker_output_token_ids = info_dict.get("thinker_output_token_ids")  # list[int]
+        prompt_embeds = embed.get("prefill")  # Tensor [P,H]
+        thinker_result = hs.get("output")  # Tensor [K,H]
+        prompt_token_ids = ids.get("prompt")  # list[int]
+        thinker_output_token_ids = ids.get("output")  # list[int]
 
         if not isinstance(prompt_embeds, torch.Tensor):
             prompt_embeds = torch.zeros(
@@ -714,7 +721,7 @@ class Qwen2_5OmniForConditionalGeneration(
         )
 
         if thinker_result.ndim == 2 and thinker_result.shape[0] > 0:
-            update_dict["thinker_reply_part"] = thinker_result[1:].detach().to("cpu").contiguous()
+            update_dict.setdefault("embed", {})["thinker_reply"] = thinker_result[1:].detach().to("cpu").contiguous()
 
         return req_input_ids, req_embeds, update_dict
 
@@ -763,18 +770,20 @@ class Qwen2_5OmniForConditionalGeneration(
             )
         return prompt_token_ids_processed, prompt_embeds
 
-    def thinker_to_talker_decode_one_step(self, input_ids, input_embeds, **info_dict):
+    def thinker_to_talker_decode_one_step(self, input_ids, input_embeds, payload: OmniPayload):
+        embed = payload.get("embed", {})
+
         update_dict = {}
         # choose step vector in priority order
         step_vec = None
-        q = info_dict.get("thinker_reply_part", None)
+        q = embed.get("thinker_reply", None)
         if isinstance(q, torch.Tensor) and q.numel() > 0:
             step_vec = q[0:1]
             new_q = q[1:].detach().to("cpu").contiguous()
-            update_dict["thinker_reply_part"] = new_q
+            update_dict.setdefault("embed", {})["thinker_reply"] = new_q
         else:
             # B) per-request provided decode vector (optional)
-            dv = info_dict.get("decode_output_prompt_embeds") if isinstance(info_dict, dict) else None
+            dv = embed.get("decode")
             if isinstance(dv, torch.Tensor) and dv.numel() > 0:
                 step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
             elif (

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -36,6 +36,7 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload, OmniPayloadMeta
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
@@ -175,15 +176,15 @@ class Qwen3OmniMoeForConditionalGeneration(
             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",
+            self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
+                ("hidden_states", "last"),
+                ("hidden_states", "trailing_text"),
+                ("embed", "tts_pad_projected"),
             }
             # Keys that need to be accumulated across streaming inputs
-            self.streaming_accumulated_keys: set[str] = {
-                "thinker_prefill_embeddings",
-                "thinker_hidden_states",
+            self.streaming_accumulated_keys: set[tuple[str, str]] = {
+                ("embed", "prefill"),
+                ("hidden_states", "output"),
             }
 
         elif self.model_stage == "code2wav":
@@ -429,8 +430,9 @@ class Qwen3OmniMoeForConditionalGeneration(
             left_context_size = []
             if runtime_additional_information is not None:
                 for info in runtime_additional_information:
-                    if "left_context_size" in info:
-                        left_context_size.append(info["left_context_size"])
+                    meta = info.get("meta", {})
+                    if "left_context_size" in meta:
+                        left_context_size.append(meta["left_context_size"])
             else:
                 logger.debug("No additional_information provided to code2wav stage.")
             audio_tensors = self.generate_audio(codes, left_context_size, seq_token_counts)
@@ -459,7 +461,7 @@ class Qwen3OmniMoeForConditionalGeneration(
             text_hidden_states, captured_layer_dict = model_outputs
             # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs.
             # These will later be projected into talker text space by the talker stage.
-            multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {}
+            multimodal_outputs: OmniPayload = captured_layer_dict if captured_layer_dict is not None else {}
             try:
                 thinker_tts_embeds = self.thinker.embed_input_ids(self.tts_tokens)  # [1,3,thinker_hidden]
                 if (
@@ -468,9 +470,10 @@ class Qwen3OmniMoeForConditionalGeneration(
                     and thinker_tts_embeds.shape[1] == 3
                 ):
                     bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1)  # 3 * [1,1,H]
-                    multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]]
-                    multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]]
-                    multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]]
+                    embed = multimodal_outputs.setdefault("embed", {})
+                    embed["tts_bos"] = [bos_eos_pad[0]]
+                    embed["tts_eos"] = [bos_eos_pad[1]]
+                    embed["tts_pad"] = [bos_eos_pad[2]]
             except Exception:
                 # Best-effort; absence will be handled by talker with fallbacks
                 pass
@@ -494,9 +497,10 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
                 logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
-            code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts]
-            multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)}
-            span_len = multimodal_outputs["code_predictor_codes"].shape[0]
+            code_predictor_codes = [info.get("codes", {}).get("audio") for info in info_dicts]
+            audio_codes = torch.cat(code_predictor_codes, dim=0)
+            multimodal_outputs: OmniPayload = {"codes": {"audio": audio_codes}}
+            span_len = audio_codes.shape[0]
             talker_hidden = talker_hidden[:span_len]
             return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs)
         elif self.model_stage == "code2wav":
@@ -640,47 +644,48 @@ class Qwen3OmniMoeForConditionalGeneration(
         """
         Postprocess the talker hidden states.
         """
-        update_dict = {}
-        update_dict["last_talker_hidden"] = hidden_states[-1, :].detach()
-        return update_dict
+        return {"hidden_states": {"last": hidden_states[-1, :].detach()}}
 
     def talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
         """
         Preprocess talker embeds. Noted that we set the MTP here.
         """
+        payload: OmniPayload = info_dict
+        meta = payload.setdefault("meta", {})
+
         # Ensure we have base embeddings when only ids are provided
         if input_embeds is None and input_ids is not None:
             input_embeds = self.talker.embed_input_ids(input_ids)
 
         span_len = input_ids.shape[0]
-        update_dict = {}
+        update_dict: OmniPayload = {}
         if span_len > 1:
             # prefill
-            input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, **info_dict)
+            input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, payload)
             code_predictor_codes = torch.zeros(
                 (input_embeds.shape[0], self.talker.num_code_groups),
                 device=self._module_device(self.talker),
                 dtype=torch.long,
             )
-            update_dict["code_predictor_codes"] = code_predictor_codes
+            update_dict.setdefault("codes", {})["audio"] = code_predictor_codes
         else:
             # decode
-            if not info_dict.get("decode_flag", False):
+            if not meta.get("decode_flag", False):
                 # Prefill already consumed the first text token via the
                 # assistant bootstrap path, so decode starts from the
                 # remaining-text boundary rather than cumulative index 0.
-                prefill_consumed_text_tokens = info_dict.get("prefill_consumed_text_tokens")
+                prefill_consumed_text_tokens = meta.get("prefill_consumed_text_tokens")
                 if prefill_consumed_text_tokens is None:
                     raise RuntimeError("Missing prefill_consumed_text_tokens for talker decode handoff.")
-                info_dict["num_processed_tokens"] = prefill_consumed_text_tokens
-                update_dict["decode_flag"] = True
+                meta["num_processed_tokens"] = prefill_consumed_text_tokens
+                update_dict.setdefault("meta", {})["decode_flag"] = True
 
             last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode(
-                input_ids, input_embeds, update_dict, **info_dict
+                input_ids, input_embeds, update_dict, payload
             )
             update_dict["mtp_inputs"] = last_talker_hidden, text_step
 
-        update_dict["num_processed_tokens"] = info_dict.get("num_processed_tokens", 0) + span_len
+        update_dict.setdefault("meta", {})["num_processed_tokens"] = meta.get("num_processed_tokens", 0) + span_len
         return input_ids, input_embeds, update_dict
 
     def talker_mtp(
@@ -737,11 +742,16 @@ class Qwen3OmniMoeForConditionalGeneration(
         self.tts_pad_embed = _proj_from_thinker(tts_pad_thinker)
         return self.tts_bos_embed, self.tts_eos_embed, self.tts_pad_embed
 
-    def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
-        # Containers to return per-request updates (e.g., code_predictor_hidden_per_request)
-        update_dict: dict[str, dict] = {}
+    def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, payload: OmniPayload):
+        hs: HiddenStates = payload.get("hidden_states", {})
+        embed: Embeddings = payload.get("embed", {})
+        ids: Ids = payload.get("ids", {})
+        meta: OmniPayloadMeta = payload.get("meta", {})
 
-        voice_type = info_dict.get("speaker")
+        # Containers to return per-request updates (e.g., code_predictor_hidden_per_request)
+        update_dict: OmniPayload = {}
+
+        voice_type = payload.get("speaker")
         if voice_type is not None and isinstance(voice_type, (list, tuple)) and len(voice_type) > 0:
             voice_type = voice_type[0]
         if not isinstance(voice_type, str) or not voice_type.strip():
@@ -749,40 +759,34 @@ class Qwen3OmniMoeForConditionalGeneration(
             voice_type = self.default_tts_text_spk_type
         else:
             voice_type = str(voice_type).lower().strip()
-        start_index = info_dict.get("num_processed_tokens", 0)
+        start_index = meta.get("num_processed_tokens", 0)
         end_index = start_index + input_embeds.shape[0]
         # Read thinker outputs for prefill
-        thinker_sequence_embeds = info_dict.get("thinker_prefill_embeddings").to(
+        thinker_sequence_embeds = embed["prefill"].to(
             device=self._module_device(self.talker), dtype=torch.bfloat16
         )  # Tensor [P,H]
-        thinker_hidden_states = info_dict.get("thinker_hidden_states").to(
+        thinker_hidden_states = hs["output"].to(
             device=self._module_device(self.talker), dtype=torch.bfloat16
         )  # Tensor [K,H]
         thinker_sequences = (
-            info_dict.get("thinker_sequences")
-            if info_dict.get("thinker_sequences") is None
-            else torch.as_tensor(info_dict.get("thinker_sequences"), device=self._module_device(self.talker))
+            ids.get("all")
+            if ids.get("all") is None
+            else torch.as_tensor(ids["all"], device=self._module_device(self.talker))
         )
         thinker_chatml_ids = (
-            info_dict.get("thinker_input_ids")
-            if info_dict.get("thinker_input_ids") is None
-            else torch.as_tensor(info_dict.get("thinker_input_ids"), device=self._module_device(self.talker))
+            ids.get("prompt")
+            if ids.get("prompt") is None
+            else torch.as_tensor(ids["prompt"], device=self._module_device(self.talker))
         )
 
-        tts_bos_thinker = info_dict.get("tts_bos_embed").to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
-        )
-        tts_eos_thinker = info_dict.get("tts_eos_embed").to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
-        )
-        tts_pad_thinker = info_dict.get("tts_pad_embed").to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
-        )
+        tts_bos_thinker = embed["tts_bos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
+        tts_eos_thinker = embed["tts_eos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
+        tts_pad_thinker = embed["tts_pad"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
 
         if thinker_sequence_embeds is None or thinker_hidden_states is None:
             raise ValueError(
                 "additional_information_by_req_id must include "
-                "'thinker_prefill_embeddings' and 'thinker_hidden_states' for talker prefill."
+                "'embed.prefill' and 'hidden_states.output' for talker prefill."
             )
 
         # Normalize to tensors
@@ -837,7 +841,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.setdefault("hidden_states", {})["trailing_text"] = rem_tail.detach()
             # Also persist projected tts_pad for decode fallback if needed
             if isinstance(tts_pad_thinker, torch.Tensor):
                 pad_in = tts_pad_thinker
@@ -846,27 +850,27 @@ 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.setdefault("embed", {})["tts_pad_projected"] = pad_proj.detach()
         except Exception:
             pass
-        update_dict["prefill_consumed_text_tokens"] = 1
-        self._talker_cache_thinker_decode_embeds(info_dict, update_dict)
+        update_dict.setdefault("meta", {})["prefill_consumed_text_tokens"] = 1
+        self._talker_cache_thinker_decode_embeds(embed, update_dict)
 
         return req_input_ids[start_index:end_index], req_embeds[start_index:end_index], update_dict
 
     def _talker_cache_thinker_decode_embeds(
         self,
-        info_dict: dict[str, Any],
-        update_dict: dict[str, Any],
+        embed: Embeddings,
+        update_dict: OmniPayload,
     ) -> None:
         """
         Cache thinker embeds for decode stage.
         """
-        thinker_decode_embeds = info_dict.get("thinker_decode_embeddings", None)
+        thinker_decode_embeds = embed.get("decode", None)
         if thinker_decode_embeds is not None:
-            cached_thinker_decode_embeds = info_dict.get("cached_thinker_decode_embeddings", None)
+            cached_thinker_decode_embeds = embed.get("cached_decode", None)
             if cached_thinker_decode_embeds is None:
-                update_dict["cached_thinker_decode_embeddings"] = thinker_decode_embeds
+                update_dict.setdefault("embed", {})["cached_decode"] = thinker_decode_embeds
             else:
                 cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(
                     device=self._module_device(self.talker), dtype=torch.bfloat16
@@ -874,10 +878,10 @@ class Qwen3OmniMoeForConditionalGeneration(
                 thinker_decode_embeds = thinker_decode_embeds.to(
                     device=self._module_device(self.talker), dtype=torch.bfloat16
                 )
-                update_dict["cached_thinker_decode_embeddings"] = torch.cat(
+                update_dict.setdefault("embed", {})["cached_decode"] = torch.cat(
                     [cached_thinker_decode_embeds, thinker_decode_embeds], dim=0
                 )
-        update_dict["thinker_decode_embeddings"] = None
+        update_dict.setdefault("embed", {})["decode"] = None
 
     def _thinker_to_talker_prefill(
         self,
@@ -964,23 +968,27 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     def _thinker_decode_to_talker_decode(
         self,
-        info_dict: dict,
+        payload: OmniPayload,
         device: torch.device,
         update_dict,
     ):
         """
-        Project thinker outputs to talker inputs during prefill stage.
+        Project thinker outputs to talker inputs during decode stage.
         Returns:
-            (input_ids, input_embeds) for talker
+            text_step embedding for talker
         """
-        cached_thinker_decode_embeds = info_dict.get("cached_thinker_decode_embeddings", None)
-        thinker_decode_embed = info_dict.get("thinker_decode_embeddings", None)
-        start_index = info_dict.get("num_processed_tokens", 0)
-        thinker_output_token_ids = info_dict.get("thinker_output_token_ids", [])
+        embed = payload.get("embed", {})
+        meta = payload.get("meta", {})
+        ids = payload.get("ids", {})
+
+        cached_thinker_decode_embeds = embed.get("cached_decode", None)
+        thinker_decode_embed = embed.get("decode", None)
+        start_index = meta.get("num_processed_tokens", 0)
+        thinker_output_token_ids = ids.get("output", [])
         if start_index >= len(thinker_output_token_ids) - 1:
-            if info_dict.get("finished_flag"):
+            if meta.get("finished"):
                 return self.tts_pad_embed.to(device)
-            update_dict["finished_flag"] = True
+            update_dict.setdefault("meta", {})["finished"] = True
             return self.tts_eos_embed.to(device)
 
         if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]:
@@ -989,25 +997,26 @@ class Qwen3OmniMoeForConditionalGeneration(
             if thinker_decode_embed is not None:
                 thinker_decode_embed = thinker_decode_embed.to(device)
                 cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0)
-                update_dict["cached_thinker_decode_embeddings"] = cached_thinker_decode_embeds
+                update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds
         else:
             thinker_embed = thinker_decode_embed
             if thinker_embed.device != device:
                 thinker_embed = thinker_embed.to(device)
-
-        update_dict["thinker_decode_embeddings"] = None
+        update_dict.setdefault("embed", {})["decode"] = None
         return self.talker.text_projection(thinker_embed).to(device)
 
     def talker_preprocess_decode(
-        self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: dict, **info_dict: dict
+        self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: OmniPayload, payload: OmniPayload
     ):
+        hs = payload.get("hidden_states", {})
+
         last_talker_hidden = None
         text_step = None
         try:
             if self.vllm_config.model_config.async_chunk:
-                text_step = self._thinker_decode_to_talker_decode(info_dict, input_ids.device, update_dict)
+                text_step = self._thinker_decode_to_talker_decode(payload, input_ids.device, update_dict)
             else:
-                q_tail = info_dict.get("trailing_text_hidden", None)
+                q_tail = hs.get("trailing_text", None)
                 if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
                     use_vec = q_tail[0:1, :]
                     new_q_tail = (
@@ -1016,11 +1025,11 @@ class Qwen3OmniMoeForConditionalGeneration(
                         else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
                     )
                     text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype)
-                    update_dict["trailing_text_hidden"] = new_q_tail
+                    update_dict.setdefault("hidden_states", {})["trailing_text"] = new_q_tail
                 else:
                     text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
 
-            last_talker_hidden_tensor = info_dict.get("last_talker_hidden")
+            last_talker_hidden_tensor = hs.get("last")
             if last_talker_hidden_tensor is not None:
                 last_talker_hidden = last_talker_hidden_tensor.to(input_embeds.device, dtype=input_embeds.dtype)
                 last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:])  # [1, hidden_size]

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -545,7 +545,9 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
 
             if captured_hidden_states is not None and capture_set is not None:
                 if layer_idx in capture_set:
-                    captured_hidden_states[str(layer_idx)] = hidden_states.clone().view(-1, hidden_states.shape[-1])
+                    hs = captured_hidden_states.setdefault("hidden_states", {})
+                    layers = hs.setdefault("layers", {})
+                    layers[layer_idx] = hidden_states.clone().view(-1, hidden_states.shape[-1])
 
             hidden_states, residual = layer(
                 positions,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -234,9 +234,10 @@ class Qwen3TTSCode2Wav(nn.Module):
             for i, info in enumerate(runtime_additional_information):
                 if i >= len(left_context_size):
                     break
-                if "left_context_size" in info:
+                meta = info.get("meta", {})
+                if "left_context_size" in meta:
                     # left_context_size may come through serialization as an int, [int], or tensor([int]).
-                    value = info["left_context_size"]
+                    value = meta["left_context_size"]
                     if isinstance(value, list):
                         value = value[0] if value else 0
                     if isinstance(value, torch.Tensor):

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -26,6 +26,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer,
 from vllm.multimodal.audio import AudioResampler
 from vllm.sequence import IntermediateTensors
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.utils.audio import mel_filter_bank
 from vllm_omni.utils.voice_cache import VoiceEmbeddingCache
@@ -343,7 +344,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self.mtp_hidden_size = int(self.talker_config.hidden_size)
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
-        self.talker_mtp_output_key = "audio_codes"
+        self.talker_mtp_output_key = ("codes", "audio")
 
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
@@ -402,11 +403,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         # Keys that should stay on GPU in model_intermediate_buffer to avoid
         # CPU-to-GPU round-trips on every decode step.
-        self.gpu_resident_buffer_keys: set[str] = {
-            "audio_codes",
-            "last_talker_hidden",
-            "tts_pad_embed",
-            "tailing_text_hidden",
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
+            ("codes", "audio"),
+            ("hidden_states", "last"),
+            ("embed", "tts_pad"),
+            ("hidden_states", "trailing_text"),
         }
 
         # Tokenizer for prompt building.
@@ -475,18 +476,20 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         for info in info_dicts:
             if not isinstance(info, dict):
                 continue
-            ac = info.get("audio_codes")
+            codes = info.get("codes", {})
+            meta = info.get("meta", {})
+            ac = codes.get("audio")
             if isinstance(ac, torch.Tensor):
                 audio_codes_list.append(ac)
-                cs = info.get("codec_streaming")
+                cs = meta.get("codec_streaming")
                 if isinstance(cs, bool):
                     codec_streaming_list.append(
                         torch.full((int(ac.shape[0]),), int(cs), dtype=torch.int8, device=ac.device)
                     )
-            ref_code = info.get("ref_code")
+            ref_code = codes.get("ref")
             if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
                 ref_code_tensor = ref_code
-            ref_len = info.get("ref_code_len")
+            ref_len = meta.get("ref_code_len")
             if ref_len is None:
                 continue
             if isinstance(ref_len, torch.Tensor):
@@ -511,13 +514,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         audio_codes = torch.cat(audio_codes_list, dim=0)
         span_len = int(audio_codes.shape[0])
         hidden = hidden[:span_len]
-        mm: dict[str, torch.Tensor] = {"audio_codes": audio_codes}
+        mm: OmniPayload = {"codes": {"audio": audio_codes}}
         if ref_code_len_list:
-            mm["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
+            mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
         if ref_code_tensor is not None:
-            mm["ref_code"] = [ref_code_tensor]
+            mm.setdefault("codes", {})["ref"] = [ref_code_tensor]
         if codec_streaming_list:
-            mm["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
+            mm.setdefault("meta", {})["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
 
     # -------------------- preprocess / postprocess --------------------
@@ -536,6 +539,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 merged.setdefault(k, v)
             info_dict = merged
 
+        payload: OmniPayload = info_dict
+        embed = payload.get("embed", {})
+        hs = payload.get("hidden_states", {})
+        meta = payload.get("meta", {})
+
         span_len = int(input_ids.shape[0])
         if span_len <= 0:
             return input_ids, input_embeds if input_embeds is not None else self.embed_input_ids(input_ids), {}
@@ -545,7 +553,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.")
 
         task_type = (info_dict.get("task_type") or ["CustomVoice"])[0]
-        codec_streaming_val = info_dict.get("codec_streaming")
+        codec_streaming_val = meta.get("codec_streaming")
         if isinstance(codec_streaming_val, list):
             codec_streaming_raw = codec_streaming_val[0] if codec_streaming_val else None
         else:
@@ -557,8 +565,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         if span_len > 1:
             # Prefill (prompt embeddings)
-            prompt_embeds_cpu = info_dict.get("talker_prompt_embeds")
-            tts_pad_embed_cpu = info_dict.get("tts_pad_embed")
+            prompt_embeds_cpu = embed.get("prefill")
+            tts_pad_embed_cpu = embed.get("tts_pad")
             tts_pad_embed = None
             if isinstance(tts_pad_embed_cpu, torch.Tensor) and tts_pad_embed_cpu.numel() > 0:
                 tts_pad_embed = tts_pad_embed_cpu.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
@@ -573,17 +581,18 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 # Store full prompt embeddings on CPU (large, prefill-only).
                 # tailing_text_hidden and tts_pad_embed stay on GPU (gpu_resident_buffer_keys).
                 prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous()
-                info_update: dict[str, Any] = {
-                    "talker_prompt_embeds": prompt_embeds_cpu,
-                    "tailing_text_hidden": tailing_text_hidden.detach(),
-                    "tts_pad_embed": tts_pad_embed.detach(),
-                    "talker_prefill_offset": 0,
-                    "codec_streaming": codec_streaming,
+                info_update: OmniPayload = {
+                    "embed": {
+                        "prefill": prompt_embeds_cpu,
+                        "tts_pad": tts_pad_embed.detach(),
+                    },
+                    "hidden_states": {"trailing_text": tailing_text_hidden.detach()},
+                    "meta": {"talker_prefill_offset": 0, "codec_streaming": codec_streaming},
                 }
                 if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-                    info_update["ref_code"] = ref_code.detach().to("cpu").contiguous()
+                    info_update.setdefault("codes", {})["ref"] = ref_code.detach().to("cpu").contiguous()
                 if ref_code_len is not None:
-                    info_update["ref_code_len"] = int(ref_code_len)
+                    info_update["meta"]["ref_code_len"] = int(ref_code_len)
                 # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed.
                 # This preserves placeholder/embedding alignment.
                 offset = 0
@@ -595,12 +604,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
                     take = torch.cat([take, pad_rows], dim=0)
                 prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
-                info_update["talker_prefill_offset"] = int(offset + span_len)
+                info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
             else:
                 # Subsequent prefill chunk: slice from stored embeddings at running offset.
                 if tts_pad_embed is None:
                     raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must initialize it.")
-                offset = int(info_dict.get("talker_prefill_offset", 0) or 0)
+                offset = int(meta.get("talker_prefill_offset", 0) or 0)
                 if offset < 0:
                     offset = 0
                 s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
@@ -611,8 +620,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
                     take = torch.cat([take, pad_rows], dim=0)
                 prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
-                info_update = {"talker_prefill_offset": int(offset + span_len)}
-                info_update["codec_streaming"] = codec_streaming
+                info_update = {
+                    "meta": {"talker_prefill_offset": int(offset + span_len), "codec_streaming": codec_streaming}
+                }
 
             # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping.
             input_ids_out = input_ids.clone()
@@ -623,18 +633,18 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 device=input_ids.device,
                 dtype=torch.long,
             )
-            info_update["audio_codes"] = zeros
+            info_update.setdefault("codes", {})["audio"] = zeros
             return input_ids_out, prompt_embeds, info_update
 
         # Decode: span_len == 1
         # Pop one text-step vector from tailing_text_hidden queue.
         # These tensors stay on GPU via gpu_resident_buffer_keys - .to() is a no-op.
-        tts_pad_embed_buf = info_dict.get("tts_pad_embed")
+        tts_pad_embed_buf = embed.get("tts_pad")
         if not isinstance(tts_pad_embed_buf, torch.Tensor):
             raise RuntimeError("Missing `tts_pad_embed` in additional_information; prefill must run first.")
         tts_pad_embed = tts_pad_embed_buf.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
 
-        tail = info_dict.get("tailing_text_hidden")
+        tail = hs.get("trailing_text")
         if isinstance(tail, torch.Tensor) and tail.ndim == 2 and tail.shape[0] > 0:
             text_step = tail[:1].to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
             new_tail = tail[1:] if tail.shape[0] > 1 else tail[:0]
@@ -642,9 +652,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             text_step = tts_pad_embed
             new_tail = tail if isinstance(tail, torch.Tensor) else torch.empty((0, tts_pad_embed.shape[-1]))
 
-        last_hidden = info_dict.get("last_talker_hidden")
+        last_hidden = hs.get("last")
         if not isinstance(last_hidden, torch.Tensor):
-            raise RuntimeError("Missing `last_talker_hidden` in additional_information; postprocess must run.")
+            raise RuntimeError("Missing hidden_states['last'] in additional_information; postprocess must run.")
         past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
 
         # Use OmniGPUModelRunner talker_mtp fast-path for residual codebooks and per-step inputs_embeds update.
@@ -654,9 +664,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         inputs_embeds_out = last_id_hidden.reshape(1, -1)
 
         info_update = {
-            "tailing_text_hidden": new_tail,
+            "hidden_states": {"trailing_text": new_tail},
             "mtp_inputs": (past_hidden, text_step),
-            "codec_streaming": codec_streaming,
+            "meta": {"codec_streaming": codec_streaming},
         }
         return input_ids, inputs_embeds_out, info_update
 
@@ -666,7 +676,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         if hidden_states.numel() == 0:
             return {}
         last = hidden_states[-1, :].detach()
-        return {"last_talker_hidden": last}
+        return {"hidden_states": {"last": last}}
 
     # -------------------- prompt construction helpers --------------------
 

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -82,7 +82,7 @@ def text2flow(
         output_ids = _ensure_list(output.token_ids)
         prefix_ids = _ensure_list(source_output.prompt_token_ids)
         additional_info = dict(multi_modal_data)
-        additional_info["prefix_ids"] = prefix_ids
+        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
         engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
     return engine_inputs
 
@@ -181,8 +181,8 @@ def talker2code2wav_async_chunk(
             if not finished:
                 return None
             payload: dict[str, Any] = {
-                "code_predictor_codes": [],
-                "finished": torch.tensor(True, dtype=torch.bool),
+                "codes": {"audio": []},
+                "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
             }
             if not state.get("sent_prompt", False):
                 payload.update(state.get("prompt_payload", {}))
@@ -193,8 +193,8 @@ def talker2code2wav_async_chunk(
         emitted_token_len = int(state.get("emitted_token_len", 0))
         if finished and length <= emitted_token_len:
             payload = {
-                "code_predictor_codes": [],
-                "finished": torch.tensor(True, dtype=torch.bool),
+                "codes": {"audio": []},
+                "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
             }
             if not state.get("sent_prompt", False):
                 payload.update(state.get("prompt_payload", {}))
@@ -225,12 +225,12 @@ def talker2code2wav_async_chunk(
             code_predictor_codes = [int(frame[0]) for frame in token_frames[:prefix_len]]
 
         payload = {
-            "code_predictor_codes": code_predictor_codes,
+            "codes": {"audio": code_predictor_codes},
+            "meta": {"finished": torch.tensor(finished, dtype=torch.bool)},
             "token_offset": token_offset,
             "left_context_size": token_offset,
             "req_id": [request_id],
             "stream_finished": torch.tensor(finished, dtype=torch.bool),
-            "finished": torch.tensor(finished, dtype=torch.bool),
         }
         if not state.get("sent_prompt", False):
             payload.update(state.get("prompt_payload", {}))

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -110,8 +110,8 @@ def slow_ar_to_dac_decoder_async_chunk(
     if length <= 0:
         if finished:
             return {
-                "code_predictor_codes": [],
-                "finished": True,
+                "codes": {"audio": []},
+                "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
             }
         return None
 
@@ -143,7 +143,6 @@ def slow_ar_to_dac_decoder_async_chunk(
     code_predictor_codes = stacked_frames.transpose(0, 1).reshape(-1).tolist()
 
     return {
-        "code_predictor_codes": code_predictor_codes,
-        "left_context_size": left_context_size,
-        "finished": finished,
+        "codes": {"audio": code_predictor_codes},
+        "meta": {"left_context_size": left_context_size, "finished": torch.tensor(finished, dtype=torch.bool)},
     }

vllm_omni/model_executor/stage_input_processors/glm_image.py

@@ -290,9 +290,8 @@ def ar2diffusion(
 
         if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
             mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict):
-                if mm_output.get("prior_token_image_ids") is not None:
-                    is_i2i = True
+            if isinstance(mm_output, dict) and mm_output.get("ids", {}).get("prior_image") is not None:
+                is_i2i = True
         _dt_mode = (time.perf_counter() - _t_mode) * 1000
 
         # Parse and upsample prior tokens
@@ -329,7 +328,7 @@ def ar2diffusion(
         if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
             mm_output = ar_output.multimodal_output
             if isinstance(mm_output, dict):
-                raw_prior_image_ids = mm_output.get("prior_token_image_ids")
+                raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                 if raw_prior_image_ids is not None:
                     # Handle different formats:
                     # 1. Single tensor -> wrap in list
@@ -355,7 +354,7 @@ def ar2diffusion(
                 mm_output = output.multimodal_output
                 logger.debug(f"[ar2diffusion] Request {i}: found multimodal_output on CompletionOutput (fallback)")
                 if isinstance(mm_output, dict):
-                    raw_prior_image_ids = mm_output.get("prior_token_image_ids")
+                    raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                     if raw_prior_image_ids is not None:
                         if isinstance(raw_prior_image_ids, torch.Tensor):
                             prior_token_image_ids = [raw_prior_image_ids]

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -55,7 +55,7 @@ def prepend_and_flatten_colmajor(x: torch.Tensor, pad_vec: torch.Tensor) -> torc
 
 def _make_finished_sentinel() -> dict[str, Any]:
     """Return a minimal payload with finished=True so Stage-1 can end the request."""
-    return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+    return {"codes": {"audio": []}, "meta": {"finished": torch.tensor(True, dtype=torch.bool)}}
 
 
 def _flush_remaining_codes(
@@ -79,12 +79,14 @@ def _flush_remaining_codes(
     flat_codes = torch.tensor(accumulated[-end_index:]).reshape(-1).tolist()
 
     return {
-        "code_predictor_codes": flat_codes,
-        "left_context_size": left_ctx_frames,
-        "codec_chunk_frames": chunk_size,
-        "codec_left_context_frames": left_context_size,
-        "code_flat_numel": len(flat_codes),
-        "finished": torch.tensor(True, dtype=torch.bool),
+        "codes": {"audio": flat_codes},
+        "meta": {
+            "left_context_size": left_ctx_frames,
+            "codec_chunk_frames": chunk_size,
+            "codec_left_context_frames": left_context_size,
+            "code_flat_numel": len(flat_codes),
+            "finished": torch.tensor(True, dtype=torch.bool),
+        },
     }
 
 
@@ -122,7 +124,6 @@ def llm2code2wav_async_chunk(
     Accumulates codes in connector per request_id,
     returns payload only when chunk_size is full or request is finished; returns None when waiting.
     """
-
     connector = getattr(transfer_manager, "connector", None)
     raw_cfg = getattr(connector, "config", {}) or {}
     cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
@@ -131,14 +132,14 @@ def llm2code2wav_async_chunk(
 
     request_id = getattr(request, "external_req_id", None)
 
-    codes = pooling_output.get("code_predictor_codes")
-
-    if _is_codes_empty(codes):
+    po_codes = pooling_output.get("codes", {})
+    if "audio" not in po_codes:
         if is_finished:
             return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
         return None
 
-    code_tensor = _to_code_tensor(codes)
+    code_predictor_codes = po_codes["audio"]
+    code_tensor = _to_code_tensor(code_predictor_codes)
     if code_tensor is None:
         if is_finished:
             return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size)
@@ -167,12 +168,16 @@ def llm2code2wav_async_chunk(
     flat_codes = torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist()
 
     return {
-        "code_predictor_codes": flat_codes,
-        "left_context_size": left_ctx_frames,
-        "codec_chunk_frames": chunk_size,
-        "codec_left_context_frames": left_context_size,
-        "code_flat_numel": len(flat_codes),
-        "finished": torch.tensor(is_finished, dtype=torch.bool),
+        "codes": {
+            "audio": flat_codes,
+        },
+        "meta": {
+            "left_context_size": left_ctx_frames,
+            "codec_chunk_frames": chunk_size,
+            "codec_left_context_frames": left_context_size,
+            "code_flat_numel": len(flat_codes),
+            "finished": torch.tensor(is_finished, dtype=torch.bool),
+        },
     }
 
 
@@ -218,8 +223,11 @@ def llm2code2wav(
 
         # Extract codec codes from talker output
         # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        if "code_predictor_codes" in output.multimodal_output:
-            codec_codes = output.multimodal_output["code_predictor_codes"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
+        mm_hs = mm.get("hidden_states", {})
+        if "audio" in mm_codes:
+            codec_codes = mm_codes["audio"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
             is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
             non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
             if len(non_zero_indices) == 0:
@@ -233,7 +241,7 @@ def llm2code2wav(
             else:
                 if len(non_zero_indices) < codec_codes.shape[0]:
                     codec_codes = codec_codes[non_zero_indices]
-        elif "latent" in output.multimodal_output and "code_predictor_codes" not in output.multimodal_output:
+        elif "output" in mm_hs and "audio" not in mm_codes:
             codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
         else:
             raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,6 +1,7 @@
 import torch
 from vllm.inputs import TextPrompt
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 TALKER_CODEC_PAD_TOKEN_ID = 8292
@@ -34,15 +35,19 @@ def thinker2talker(
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        latent = output.multimodal_output["latent"]
+        mm: OmniPayload = output.multimodal_output
+        latent = mm["latent"]
         thinker_hidden_states = latent.clone().detach().to(latent.device)
         additional_information = {
-            "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
-            "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
-            "prompt_token_ids": prompt_token_ids,
-            "thinker_output_token_ids": thinker_output_ids,
-            "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
-            "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
+            "hidden_states": {
+                "output": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
+                "output_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
+            },
+            "embed": {
+                "prefill": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
+                "prefill_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
+            },
+            "ids": {"prompt": prompt_token_ids, "output": thinker_output_ids},
         }
         talker_inputs.append(
             OmniTokensPrompt(

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -11,6 +11,7 @@ import torch
 from vllm.inputs import TextPrompt
 from vllm.platforms import current_platform
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
@@ -27,15 +28,16 @@ _EMBED_LAYER_KEY = "0"
 _HIDDEN_LAYER_KEY = "24"
 
 
-def _compute_talker_prompt_ids_length(info, device: torch.device | str = "cuda") -> int:
+def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int:
     im_start_token_id = 151644
     system_token_id = 8948
     user_token_id = 872
     assistant_token_id = 77091
 
-    thinker_sequences = torch.tensor(info["thinker_sequences"], dtype=torch.long, device=device).unsqueeze(0)  # [1, T]
+    ids = info.get("ids", {})
+    thinker_sequences = torch.tensor(ids["all"], dtype=torch.long, device=device).unsqueeze(0)  # [1, T]
 
-    input_ids = torch.tensor(info["thinker_input_ids"], dtype=torch.long, device=device).unsqueeze(0)  # [1, T]
+    input_ids = torch.tensor(ids["prompt"], dtype=torch.long, device=device).unsqueeze(0)  # [1, T]
 
     im_start_indexes = torch.cat(
         [
@@ -127,8 +129,9 @@ def _merge_pd_embeddings(
     where overlap = P + D - expected_total.
     """
     try:
-        p_emb = prefill_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float)
-        p_hid = prefill_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float)
+        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, dict) else type(prefill_mm).__name__
         logger.error(
@@ -175,9 +178,9 @@ def _resolve_tts_token_embedding(
     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(key)
+    val = thinker_mm.get("embed", {}).get(key)
     if val is None and prefill_mm is not None:
-        val = prefill_mm.get(key)
+        val = prefill_mm.get("embed", {}).get(key)
     return val.detach().to(device=device, dtype=torch.float) if val is not None else None
 
 
@@ -297,23 +300,29 @@ def thinker2talker_async_chunk(
 
     request_id = request.external_req_id
     chunk_id = transfer_manager.put_req_chunk[request_id]
+    thinker_hs = pooling_output.get("hidden_states", {})
+    thinker_layers = thinker_hs.get("layers", {})
+    thinker_embed = pooling_output.get("embed", {})
+
     if chunk_id == 0:
         all_token_ids = request.all_token_ids  # prefill + decode
         prompt_token_ids = request.prompt_token_ids
         # Convert ConstantList to regular list for OmniSerializer serialization
         all_token_ids = _ensure_list(all_token_ids)
         prompt_token_ids = _ensure_list(prompt_token_ids)
-        talker_additional_info = {
-            "thinker_prefill_embeddings": pooling_output.get(_EMBED_LAYER_KEY).detach().cpu(),
-            "thinker_hidden_states": pooling_output.get(_HIDDEN_LAYER_KEY).detach().cpu(),
-            "thinker_sequences": all_token_ids,
-            "thinker_input_ids": prompt_token_ids,
-            # Provide thinker-side TTS token embeddings for talker projection
-            "tts_bos_embed": pooling_output.get("tts_bos_embed").detach().cpu(),
-            "tts_eos_embed": pooling_output.get("tts_eos_embed").detach().cpu(),
-            "tts_pad_embed": pooling_output.get("tts_pad_embed").detach().cpu(),
-            "finished": torch.tensor(is_finished, dtype=torch.bool),
+        payload: OmniPayload = {
+            "embed": {
+                "prefill": thinker_layers[int(_EMBED_LAYER_KEY)].detach().cpu(),
+                # Provide thinker-side TTS token embeddings for talker projection
+                "tts_bos": thinker_embed["tts_bos"].detach().cpu(),
+                "tts_eos": thinker_embed["tts_eos"].detach().cpu(),
+                "tts_pad": thinker_embed["tts_pad"].detach().cpu(),
+            },
+            "hidden_states": {"output": thinker_layers[int(_HIDDEN_LAYER_KEY)].detach().cpu()},
+            "ids": {"all": all_token_ids, "prompt": prompt_token_ids},
+            "meta": {"finished": torch.tensor(is_finished, dtype=torch.bool)},
         }
+        talker_additional_info = payload
         speaker = extract_speaker_from_request(request)
         if speaker is not None:
             talker_additional_info["speaker"] = speaker
@@ -326,15 +335,18 @@ def thinker2talker_async_chunk(
                 return None
         else:
             save_payload = transfer_manager.request_payload.pop(request_id)
-            talker_additional_info["thinker_prefill_embeddings"] = torch.cat(
+            talker_additional_info["embed"]["prefill"] = torch.cat(
                 (
-                    save_payload.get("thinker_prefill_embeddings"),
-                    talker_additional_info.get("thinker_prefill_embeddings"),
+                    save_payload.get("embed", {}).get("prefill"),
+                    talker_additional_info.get("embed", {}).get("prefill"),
                 ),
                 dim=0,
             )
-            talker_additional_info["thinker_hidden_states"] = torch.cat(
-                (save_payload.get("thinker_hidden_states"), talker_additional_info.get("thinker_hidden_states")),
+            talker_additional_info["hidden_states"]["output"] = torch.cat(
+                (
+                    save_payload.get("hidden_states", {}).get("output"),
+                    talker_additional_info.get("hidden_states", {}).get("output"),
+                ),
                 dim=0,
             )
     else:
@@ -342,8 +354,8 @@ def thinker2talker_async_chunk(
         # Convert ConstantList to regular list for OmniSerializer serialization
         output_token_ids = _ensure_list(output_token_ids)
 
-        talker_additional_info = {
-            "finished": torch.tensor(is_finished, dtype=torch.bool),
+        talker_additional_info: OmniPayload = {
+            "meta": {"finished": torch.tensor(is_finished, dtype=torch.bool)},
         }
         speaker = extract_speaker_from_request(request)
         if speaker is not None:
@@ -353,14 +365,13 @@ def thinker2talker_async_chunk(
             talker_additional_info["language"] = language
 
         if output_token_ids:
-            talker_additional_info["override_keys"] = ["thinker_decode_embeddings", "thinker_output_token_ids"]
-            talker_additional_info["thinker_decode_embeddings"] = pooling_output.get(_EMBED_LAYER_KEY).detach().cpu()
-            talker_additional_info["thinker_output_token_ids"] = output_token_ids
+            talker_additional_info["meta"]["override_keys"] = [("embed", "decode"), ("ids", "output")]
+            talker_additional_info["embed"] = {"decode": thinker_layers[int(_EMBED_LAYER_KEY)].detach().cpu()}
+            talker_additional_info["ids"] = {"output": output_token_ids}
         else:
             # When prefilling a chunked thinker, thinker_hidden_states needs to be updated.
-            talker_additional_info["thinker_prefill_embeddings"] = pooling_output.get("0").detach().cpu()
-            talker_additional_info["thinker_hidden_states"] = pooling_output.get("24").detach().cpu()
-
+            talker_additional_info["embed"] = {"prefill": thinker_layers[0].detach().cpu()}
+            talker_additional_info["hidden_states"] = {"output": thinker_layers[24].detach().cpu()}
     return talker_additional_info
 
 
@@ -419,14 +430,12 @@ def thinker2talker(
         else:
             thinker_sequences = prompt_token_ids + output_ids
             thinker_input_ids = prompt_token_ids
-        # For streaming input, just send incremental prefill and hidden states tensor to talker
-        # Equally applicable to non-streaming cases.
         new_seq_length = len(prompt_token_ids + output_ids) - 1
-        thinker_mm = output.multimodal_output
-        # Full thinker embedding sequence for the talker: single thinker engine in the
-        # non-PD path; after optional merge with prefill-side tensors in PD mode.
-        thinker_emb = thinker_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-        thinker_hid = thinker_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float)[-new_seq_length:]
+        thinker_mm: OmniPayload = output.multimodal_output
+        mm_hs = thinker_mm.get("hidden_states", {})
+        mm_layers = mm_hs.get("layers", {})
+        thinker_emb = mm_layers[int(_EMBED_LAYER_KEY)].detach().to(device=device, dtype=torch.float)[-new_seq_length:]
+        thinker_hid = mm_layers[int(_HIDDEN_LAYER_KEY)].detach().to(device=device, dtype=torch.float)[-new_seq_length:]
 
         prefill_mm: dict[str, Any] | None = None
         if prefill_stage is not None:
@@ -441,22 +450,28 @@ def thinker2talker(
             except Exception as exc:
                 logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
 
-        info = {
-            "thinker_prefill_embeddings": thinker_emb,
-            "thinker_hidden_states": thinker_hid,
-            "thinker_sequences": thinker_sequences,  # the thinker_sequences is the whole ids
-            "thinker_input_ids": thinker_input_ids,
-            # Provide thinker-side TTS token embeddings for talker projection
-            "tts_bos_embed": _resolve_tts_token_embedding(
-                "tts_bos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-            ),
-            "tts_eos_embed": _resolve_tts_token_embedding(
-                "tts_eos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-            ),
-            "tts_pad_embed": _resolve_tts_token_embedding(
-                "tts_pad_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-            ),
+        payload: OmniPayload = {
+            "embed": {
+                "prefill": thinker_emb,
+                "tts_bos": _resolve_tts_token_embedding(
+                    "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+                "tts_eos": _resolve_tts_token_embedding(
+                    "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+                "tts_pad": _resolve_tts_token_embedding(
+                    "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+            },
+            "hidden_states": {
+                "output": thinker_hid,
+            },
+            "ids": {
+                "all": thinker_sequences,
+                "prompt": thinker_input_ids,
+            },
         }
+        info = payload
         speaker = extract_speaker_from_prompt(prompt, index=i)
         if speaker is not None:
             info["speaker"] = speaker
@@ -464,7 +479,7 @@ def thinker2talker(
         if language is not None:
             info["language"] = language
 
-        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
+        prompt_len = _compute_talker_prompt_ids_length(payload, device=device)
 
         talker_inputs.append(
             OmniTokensPrompt(
@@ -492,7 +507,9 @@ def talker2code2wav_async_chunk(
     """
     Pooling version.
     """
-    if "code_predictor_codes" not in pooling_output:
+    talker_codes = pooling_output.get("codes", {})
+    code_predictor_codes = talker_codes.get("audio")
+    if code_predictor_codes is None:
         return None
 
     connector = getattr(transfer_manager, "connector", None)
@@ -501,8 +518,6 @@ def talker2code2wav_async_chunk(
     chunk_size_config = int(cfg.get("codec_chunk_frames", 25))
     left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
 
-    code_predictor_codes = pooling_output["code_predictor_codes"]
-
     if code_predictor_codes is None:
         return None
     if isinstance(code_predictor_codes, torch.Tensor):
@@ -527,6 +542,7 @@ def talker2code2wav_async_chunk(
     request_id = request.external_req_id
     transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
     length = len(transfer_manager.code_prompt_token_ids[request_id])
+
     chunk_length = length % chunk_size_config
     if chunk_length != 0 and not is_finished:
         return None
@@ -543,12 +559,10 @@ def talker2code2wav_async_chunk(
         .tolist()
     )
 
-    info = {
-        "code_predictor_codes": codes,
-        "left_context_size": left_context_size,
-        "finished": torch.tensor(is_finished, dtype=torch.bool),
+    return {
+        "codes": {"audio": codes},
+        "meta": {"left_context_size": left_context_size, "finished": torch.tensor(is_finished, dtype=torch.bool)},
     }
-    return info
 
 
 def talker2code2wav(
@@ -586,16 +600,11 @@ def talker2code2wav(
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
             seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
+        mm: OmniPayload = output.multimodal_output
         # Extract codec codes from talker output
         # Expected shape: [8, seq_len] (8-layer RVQ codes)
         codec_codes = (
-            output.multimodal_output["code_predictor_codes"][-seq_len:]
-            .to(torch.long)
-            .transpose(0, 1)
-            .cpu()
-            .to(torch.long)
-            .reshape(-1)
-            .tolist()
+            mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
         )  # 16, seq_len
         code2wav_inputs.append(
             OmniTokensPrompt(

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -37,8 +37,10 @@ def talker2code2wav(
             # accumulated the final code sequence.
             continue
         output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
         # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
-        audio_codes = output.multimodal_output["audio_codes"].to(torch.long)
+        audio_codes = mm_codes["audio"].to(torch.long)
         token_ids = output.token_ids
         # token_ids provides an upper bound on the newly generated codec span.
         # audio_codes may still contain zero-padded / invalid rows, so trim only
@@ -51,8 +53,8 @@ def talker2code2wav(
         audio_codes = audio_codes[valid_mask]
         if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
             audio_codes = audio_codes[-seq_len:]
-        ref_code = output.multimodal_output.get("ref_code")
-        ref_code_len = output.multimodal_output.get("ref_code_len")
+        ref_code = mm_codes.get("ref")
+        ref_code_len = mm.get("meta", {}).get("ref_code_len")
         if isinstance(ref_code_len, torch.Tensor):
             ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
         elif ref_code_len is None:
@@ -95,7 +97,7 @@ def talker2code2wav(
         codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
         additional_information: dict[str, Any] = {}
         if ref_code_len > 0:
-            additional_information["left_context_size"] = [ref_code_len]
+            additional_information["meta"] = {"left_context_size": [ref_code_len]}
         # Propagate speaker and language from the original prompt so they are
         # available as runtime_additional_information in later pipeline stages,
         # consistent with qwen3-omni and qwen2.5-omni stage input processors.
@@ -117,7 +119,7 @@ def talker2code2wav(
 
 
 def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
-    audio_codes = pooling_output.get("audio_codes")
+    audio_codes = pooling_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return None
     if audio_codes.ndim == 2:
@@ -148,7 +150,7 @@ def talker2code2wav_async_chunk(
         if frame is not None:
             codec_codes = frame.cpu().tolist()
             transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
-        ref_code = pooling_output.get("ref_code")
+        ref_code = pooling_output.get("codes", {}).get("ref")
         if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None:
             request_payload[request_id] = ref_code.to(torch.long).cpu().contiguous()
     elif not finished:
@@ -207,8 +209,8 @@ def talker2code2wav_async_chunk(
     if length <= 0:
         if finished:
             return {
-                "code_predictor_codes": [],
-                "finished": True,
+                "codes": {"audio": []},
+                "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
             }
         return None
 
@@ -254,9 +256,8 @@ def talker2code2wav_async_chunk(
     code_predictor_codes = [window_frames[f][q] for q in range(num_quantizers) for f in range(num_frames)]
 
     info: dict[str, Any] = {
-        "code_predictor_codes": code_predictor_codes,
-        "left_context_size": left_context_size,
-        "finished": finished,
+        "codes": {"audio": code_predictor_codes},
+        "meta": {"left_context_size": left_context_size, "finished": torch.tensor(finished, dtype=torch.bool)},
     }
     # Propagate speaker and language from the request so they are available
     # as runtime_additional_information in subsequent pipeline stages, consistent

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -82,8 +82,8 @@ def generator2tokenizer_async_chunk(
     if length <= 0:
         if finished:
             return {
-                "code_predictor_codes": [],
-                "finished": torch.tensor(True, dtype=torch.bool),
+                "codes": {"audio": []},
+                "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
             }
         return None
 
@@ -105,6 +105,6 @@ def generator2tokenizer_async_chunk(
     code_predictor_codes = torch.tensor(window_frames).reshape(-1).tolist()
 
     return {
-        "code_predictor_codes": [int(ctx_frames)] + [int(context_length)] + code_predictor_codes,
-        "finished": torch.tensor(finished, dtype=torch.bool),
+        "codes": {"audio": [int(ctx_frames)] + [int(context_length)] + code_predictor_codes},
+        "meta": {"finished": torch.tensor(finished, dtype=torch.bool)},
     }

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -433,14 +433,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
                 top_k=subtalker_params.get("top_k"),
                 top_p=subtalker_params.get("top_p"),
             )
-        # code_predictor_codes stays on GPU here; _update_intermediate_buffer
-        # keeps it device-resident when the key is in gpu_resident_buffer_keys.
-        # D2H is deferred to sample_tokens where hidden_states.to("cpu") already
-        # syncs the stream, avoiding a per-step cudaStreamSynchronize.
-        out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
+        # update the inputs_embeds and code_predictor_codes
+        code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
+        out_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio"))
+        if not isinstance(out_key, tuple) or len(out_key) != 2:
+            raise TypeError(f"talker_mtp_output_key must be a 2-tuple, got {type(out_key).__name__}: {out_key!r}")
         for idx, req_id in enumerate(decode_req_ids):
             req_index = self.input_batch.req_ids.index(req_id)
             start_offset = int(self.query_start_loc.cpu[req_index])
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key: code_predictor_codes[idx : idx + 1]}
+            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes_cpu[idx : idx + 1]}}
             self._merge_additional_information_update(req_id, update_dict)

vllm_omni/worker/gpu_ar_model_runner.py

@@ -37,6 +37,7 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
+from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element
@@ -854,7 +855,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             )
         # Otherwise we don't have the mm CPU data yet, so we still need to build it
         if self.omni_prefix_cache is None:
-            mm_cpu = build_mm_cpu(multimodal_outputs)
+            mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs))
 
         self._process_additional_information_updates(
             hidden_states,
@@ -909,7 +910,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                             seq_len=seq_len,
                         )
                 payload.update(mm_payload)
-            pooler_output.append(payload)
+            # Flatten nested dicts to dotted keys so pooling_output
+            # stays dict[str, torch.Tensor] for msgspec serialization.
+            pooler_output.append(flatten_payload(payload))
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
             if self.routed_experts_initialized:
                 capturer = RoutedExpertsCapturer.get_instance()

vllm_omni/worker/gpu_model_runner.py

@@ -386,8 +386,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     logger.warning_once(
                         "additional_information on request data is deprecated, use model_intermediate_buffer"
                     )
-                    payload_info = new_req_data.additional_information
-                    info_dict = deserialize_additional_information(payload_info)
+                    info_dict = deserialize_additional_information(new_req_data.additional_information)
                     if info_dict:
                         self.model_intermediate_buffer[req_id] = info_dict
                         setattr(
@@ -1082,10 +1081,16 @@ class OmniGPUModelRunner(GPUModelRunner):
                         mm_out = {k: v[req_id] for k, v in combined_multimodal_outputs.items()}
                     else:
                         mm_out = multimodal_outputs
+                    # Exclude 'hidden_states' from kwargs to avoid clash with
+                    # the positional arg. The buffer entry must be preserved
+                    # because preprocess reads hidden_states['last'] from it.
+                    # TODO: pass req_infos as a single payload arg instead of **unpacking
+                    # to avoid key collisions with positional args.
+                    postprocess_kwargs = {k: v for k, v in req_infos.items() if k != "hidden_states"}
                     update_dict = self.model.postprocess(
                         hidden_states_slice,
                         multimodal_outputs=mm_out,
-                        **req_infos,
+                        **postprocess_kwargs,
                     )
                     self._update_intermediate_buffer(req_id, update_dict)
         except Exception as e:
@@ -1380,16 +1385,15 @@ class OmniGPUModelRunner(GPUModelRunner):
                 top_k=subtalker_params.get("top_k"),
                 top_p=subtalker_params.get("top_p"),
             )
-        # code_predictor_codes stays on GPU here; _update_intermediate_buffer
-        # keeps it device-resident when the key is in gpu_resident_buffer_keys.
-        # D2H is deferred to sample_tokens where hidden_states.to("cpu") already
-        # syncs the stream, avoiding a per-step cudaStreamSynchronize.
-        out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
+        # update the inputs_embeds and code_predictor_codes
+        out_key = getattr(self.model, "talker_mtp_output_key", ("codes", "audio"))
+        if not isinstance(out_key, tuple) or len(out_key) != 2:
+            raise TypeError(f"talker_mtp_output_key must be a 2-tuple, got {type(out_key).__name__}: {out_key!r}")
         for idx, req_id in enumerate(decode_req_ids):
             req_index = self.input_batch.req_ids.index(req_id)
             start_offset = int(self.query_start_loc.cpu[req_index])
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key: code_predictor_codes[idx : idx + 1]}
+            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
             self._merge_additional_information_update(req_id, update_dict)
 
     def _model_forward(
@@ -1417,29 +1421,37 @@ class OmniGPUModelRunner(GPUModelRunner):
         self._omni_last_model_output = model_output
         return model_output
 
+    def _store_value(self, dest: dict, key: str, value: Any, gpu_keys: set) -> None:
+        if isinstance(value, torch.Tensor):
+            if key in gpu_keys:
+                dest[key] = value.detach().clone()
+            else:
+                dest[key] = value.detach().to("cpu").contiguous()
+        elif isinstance(value, list):
+            dest[key] = [
+                (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in value
+            ]
+        else:
+            dest[key] = value
+
     def _update_intermediate_buffer(self, req_id: str, upd: dict) -> None:
         if not isinstance(upd, dict) or not upd:
             return
         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()
+        # Check if the model declares keys that should stay on GPU (tuples of (type_key, qualifier))
+        gpu_keys: set[tuple[str, 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()
-            elif isinstance(v, list):
-                existing[k] = [
-                    (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v
-                ]
+            if isinstance(v, dict):
+                existing_sub = existing.setdefault(k, {})
+                for qual, val in v.items():
+                    self._store_value(existing_sub, qual, val, {q for tk, q in gpu_keys if tk == k})
             else:
-                existing[k] = v
+                self._store_value(existing, k, v, set())
         # Backward compatible: mirror to old setattr location
         setattr(req_state, "additional_information_cpu", existing)
 
@@ -1454,22 +1466,27 @@ class OmniGPUModelRunner(GPUModelRunner):
             payload_info = getattr(new_req_data, "additional_information", None)
             inc_info = deserialize_additional_information(payload_info)
             if isinstance(inc_info, dict) and inc_info:
+                accumulated_keys: set[tuple[str, str]] = set()
+                if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"):
+                    accumulated_keys = self.model.streaming_accumulated_keys
                 merged_info = dict(cached_additional_info)
                 for key, value in inc_info.items():
-                    accumulated_keys: set[str] = set()
-                    if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"):
-                        accumulated_keys = self.model.streaming_accumulated_keys
-                    if key in accumulated_keys and isinstance(value, torch.Tensor):
-                        inc_tensor = value.detach().to("cpu").contiguous()
-                        old_tensor = merged_info.get(key)
-                        if old_tensor is None:
-                            merged_info[key] = inc_tensor
-                        else:
-                            merged_info[key] = torch.cat((old_tensor, inc_tensor), dim=0)
-                        continue
-
-                    # Default for other keys: latest value.
-                    merged_info[key] = value
-                merged_info["num_processed_tokens"] = 0
+                    if isinstance(value, dict):
+                        existing_sub = merged_info.get(key)
+                        merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
+                        for sk, sv in value.items():
+                            if (key, sk) in accumulated_keys and isinstance(sv, torch.Tensor):
+                                inc_tensor = sv.detach().to("cpu").contiguous()
+                                old_tensor = merged_sub.get(sk)
+                                if old_tensor is None:
+                                    merged_sub[sk] = inc_tensor
+                                else:
+                                    merged_sub[sk] = torch.cat((old_tensor, inc_tensor), dim=0)
+                            else:
+                                merged_sub[sk] = sv
+                        merged_info[key] = merged_sub
+                    else:
+                        merged_info[key] = value
+                merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
                 self.model_intermediate_buffer[req_id] = merged_info
                 setattr(self.requests[req_id], "additional_information_cpu", merged_info)

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -365,16 +365,21 @@ class OmniConnectorModelRunnerMixin:
     #  Scheduling metadata extraction
     # ------------------------------------------------------------------ #
 
-    _SCHEDULING_METADATA_KEYS = (
-        "next_stage_prompt_len",
-        "code_predictor_codes",
-        "left_context_size",
-    )
-
     @classmethod
     def _extract_scheduling_metadata(cls, payload: dict[str, Any]) -> dict[str, Any]:
         """Extract only the fields the scheduler needs from a full payload."""
-        return {k: payload[k] for k in cls._SCHEDULING_METADATA_KEYS if k in payload}
+        extracted: dict[str, Any] = {}
+        if "next_stage_prompt_len" in payload:
+            extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"]
+        audio_codes = cls._payload_audio_codes(payload)
+        if audio_codes is not None:
+            extracted["code_predictor_codes"] = audio_codes
+        meta = payload.get("meta")
+        if isinstance(meta, dict) and "left_context_size" in meta:
+            extracted["left_context_size"] = meta["left_context_size"]
+        elif "left_context_size" in payload:
+            logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'")
+        return extracted
 
     _NON_CONSUMABLE_PAYLOAD_KEYS = {
         "finished",
@@ -396,6 +401,31 @@ class OmniConnectorModelRunnerMixin:
             return len(value) > 0
         return True
 
+    @staticmethod
+    def _payload_finished(payload: Any) -> bool:
+        if not isinstance(payload, dict):
+            return False
+        if "finished" in payload:
+            logger.warning_once("legacy flat 'finished' key in payload; expected 'meta.finished'")
+        meta = payload.get("meta")
+        if not isinstance(meta, dict) or "finished" not in meta:
+            return False
+        flag = meta["finished"]
+        if isinstance(flag, torch.Tensor):
+            return flag.numel() == 1 and bool(flag.item())
+        return bool(flag)
+
+    @staticmethod
+    def _payload_audio_codes(payload: Any) -> Any:
+        if not isinstance(payload, dict):
+            return None
+        if "code_predictor_codes" in payload:
+            logger.warning_once("legacy flat 'code_predictor_codes' key in payload; expected 'codes.audio'")
+        codes = payload.get("codes")
+        if isinstance(codes, dict):
+            return codes.get("audio")
+        return None
+
     @classmethod
     def _payload_is_consumable(cls, payload: dict[str, Any] | None) -> bool:
         """Return True when an async payload can drive a real forward step.
@@ -414,15 +444,14 @@ class OmniConnectorModelRunnerMixin:
                 return True
             return decode_embeddings.numel() > 0 and decode_embeddings.shape[0] > 0
 
-        if "code_predictor_codes" in payload:
-            code_predictor_codes = payload.get("code_predictor_codes")
-            if isinstance(code_predictor_codes, torch.Tensor):
-                return code_predictor_codes.numel() > 0
+        audio_codes = cls._payload_audio_codes(payload)
+        if audio_codes is not None:
+            if isinstance(audio_codes, torch.Tensor):
+                return audio_codes.numel() > 0
             # Codec code 0 is valid; non-empty code payloads are consumable.
-            if hasattr(code_predictor_codes, "__len__"):
-                return len(code_predictor_codes) > 0
-            else:
-                return code_predictor_codes is not None
+            if hasattr(audio_codes, "__len__"):
+                return len(audio_codes) > 0
+            return True
 
         for key, value in payload.items():
             if key in cls._NON_CONSUMABLE_PAYLOAD_KEYS:
@@ -731,20 +760,21 @@ class OmniConnectorModelRunnerMixin:
                 logger.info("[Stage-%s] send_full_payload_outputs: payload is None for %s", self._stage_id, req_id)
                 continue
             if isinstance(payload, dict):
-                code_predictor_codes = payload.get("code_predictor_codes")
-                if isinstance(code_predictor_codes, torch.Tensor):
-                    code_len = int(code_predictor_codes.numel())
-                elif hasattr(code_predictor_codes, "__len__"):
-                    code_len = len(code_predictor_codes)
+                audio_codes = self._payload_audio_codes(payload)
+                if isinstance(audio_codes, torch.Tensor):
+                    code_len = int(audio_codes.numel())
+                elif hasattr(audio_codes, "__len__"):
+                    code_len = len(audio_codes)
                 else:
                     code_len = None
+                meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {}
                 logger.info(
                     "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s",
                     self._stage_id,
                     req_id,
                     sorted(payload.keys()),
                     code_len,
-                    payload.get("left_context_size"),
+                    meta.get("left_context_size"),
                 )
 
             external_req_id = self._resolve_external_req_id(request, req_id)
@@ -1545,20 +1575,20 @@ class OmniConnectorModelRunnerMixin:
                 external_req_id,
                 connector_get_key,
                 sorted(payload_data.keys()),
-                bool(payload_data.get("finished")) if "finished" in payload_data else None,
+                self._payload_finished(payload_data),
             )
 
         self._get_req_chunk[req_id] += 1
 
         if self._async_chunk:
-            is_finished = bool(payload_data.get("finished"))
+            is_finished = self._payload_finished(payload_data)
             incoming_payload_consumable = self._payload_is_consumable(payload_data)
 
             if self._model_mode == "ar":
                 payload_data = self._accumulate_payload(external_req_id, payload_data)
                 payload_consumable = incoming_payload_consumable
             else:
-                new_ids = payload_data.get("code_predictor_codes", [])
+                new_ids = self._payload_audio_codes(payload_data) or []
                 if not new_ids and not is_finished:
                     return False
                 payload_consumable = self._payload_is_consumable(payload_data)