From 9ceca4ec85c21d1806ea69d74002e0f9e96bbca7 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Mon, 6 Jul 2026 15:08:05 +0800 Subject: [PATCH] [Bugfix] Fix full-payload mm splitting for dual hidden/scheduled batch axes (#4870 follow-up) (#4910) Signed-off-by: amy-why-3459 --- tests/worker/test_gpu_ar_model_runner.py | 46 +++++++++++++++++++ .../npu/worker/npu_ar_model_runner.py | 6 ++- vllm_omni/utils/mm_outputs.py | 44 ++++++++++++++---- vllm_omni/worker/gpu_ar_model_runner.py | 18 +++++--- 4 files changed, 96 insertions(+), 18 deletions(-) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index b56917fdb..90cf3afe1 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -480,6 +480,52 @@ def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_onl assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3]) +def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch): + """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count.""" + runner = _make_async_output_runner(engine_output_type="latent") + runner.model.omni_pooler_payload_include_hidden = False + runner._async_chunk = False + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + layers = torch.arange(96, dtype=torch.long).reshape(3, 32) + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=5, + num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1}, + ), + hidden_states=torch.randn(3, 4), + staged_hidden_states_cpu=None, + multimodal_outputs={"hidden_states": {"layers": {0: layers}}}, + req_ids_output_copy=["r1", "r2", "r3"], + req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2}, + valid_sampled_token_ids=[[101], [102], [103]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1]) + assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2]) + assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3]) + + def test_async_snapshot_payload_omits_hidden_when_model_opts_out(): runner = _make_async_output_runner() runner.model.omni_pooler_payload_include_hidden = False diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 43d2c9d19..a0abc7fa2 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -902,7 +902,8 @@ class NPUARModelRunner(OmniNPUModelRunner): ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None - seq_len = hidden_states.shape[0] + hidden_seq_len = int(hidden_states.shape[0]) + scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens) # Apply structured output bitmasks if present. if grammar_output is not None: @@ -1141,7 +1142,8 @@ class NPUARModelRunner(OmniNPUModelRunner): start=start, end=end, pass_lists_through=False, - seq_len=seq_len, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, ) payload.update(mm_payload) pooler_output.append(flatten_payload(payload)) diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 0ca0c2c3a..189cc00aa 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -106,7 +106,13 @@ def _to_cpu(value): def to_payload_element( - element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None + element: object, + idx: int, + start: int, + end: int, + pass_lists_through: bool = False, + seq_len: int | None = None, + scheduled_seq_len: int | None = None, ): """Build an mm payload element corresponding to one request index from an element containing 0 or more CPU tensors. @@ -120,21 +126,39 @@ def to_payload_element( passthrough data; this should be False in normal cases, but True if we need to avoid splitting nonempty lists prior to calling postprocess, which is the case for prefix cache. - seq_len: Optional sequence length (i.e., dim 0 of hidden states). - When set, a tensor whose first dimension equals seq_len is - sliced per request. The prefix cache passthrough also passes - the total scheduled token count here so 1D (seq_len,) metadata - that is intentionally not cached is still split per request. + seq_len: Optional hidden-aligned batch length (``hidden_states.shape[0]``). + When a tensor's first dimension equals this value, it is sliced + per request using ``start:end``. + scheduled_seq_len: Optional scheduler-aligned batch length + (``scheduler_output.total_num_scheduled_tokens``). Some full-payload + mm tensors (e.g. batched ``codes.audio`` with tail-only hidden states) + are laid out by scheduled tokens instead of the hidden tail shape. + When omitted, ``seq_len`` is reused for backward compatibility. """ - # Cached per-token tensors are merged elsewhere; here a first dim - # equal to seq_len means a per-request slice is required. - if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len: + if scheduled_seq_len is None: + scheduled_seq_len = seq_len + + # Cached per-token tensors are merged elsewhere; here a first dim equal to + # either the hidden-aligned or scheduler-aligned batch length means a + # per-request slice is required. + if isinstance(element, torch.Tensor) and ( + (seq_len is not None and element.shape[0] == seq_len) + or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len) + ): return element[start:end].contiguous() # Every other case is shared between prefix cache (passthrough data) # and running a model without prefix caching. elif isinstance(element, dict): return { - sk: to_payload_element(sv, idx, start, end, pass_lists_through=pass_lists_through, seq_len=seq_len) + sk: to_payload_element( + sv, + idx, + start, + end, + pass_lists_through=pass_lists_through, + seq_len=seq_len, + scheduled_seq_len=scheduled_seq_len, + ) for sk, sv in element.items() } elif isinstance(element, list): diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 2b44e7b45..2bf7a8af2 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -839,7 +839,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): end: int, audio_sparse_output: bool, sparse_mm_index: dict[str, int], - seq_len: int, + hidden_seq_len: int, + scheduled_seq_len: int, ) -> dict[str, object]: if combined_multimodal_outputs: return self._build_combined_prefix_cache_mm_payload( @@ -876,7 +877,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): start=start, end=end, pass_lists_through=False, - seq_len=seq_len, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, ) return mm_payload @@ -894,7 +896,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): mm_cpu: dict[str, object] | None, audio_sparse_output: bool, sparse_mm_index: dict[str, int], - seq_len: int, + hidden_seq_len: int, + scheduled_seq_len: int, ) -> dict[str, object]: payload: dict[str, object] = {} if not audio_sparse_output: @@ -920,7 +923,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): end=end, audio_sparse_output=audio_sparse_output, sparse_mm_index=sparse_mm_index, - seq_len=seq_len, + hidden_seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, ) payload.update(mm_payload) return payload @@ -1731,7 +1735,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # The actual multimodal wire transport uses multimodal_outputs instead. pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: - mm_seq_len = int(scheduler_output.total_num_scheduled_tokens) + hidden_seq_len = int(hidden_states.shape[0]) + scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens) mm_cpu = None if self.omni_prefix_cache is not None: ( @@ -1787,7 +1792,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): mm_cpu=mm_cpu, audio_sparse_output=audio_sparse_output, sparse_mm_index=sparse_mm_index, - seq_len=mm_seq_len, + hidden_seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, ) pooler_output.append(flatten_payload(payload))