diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index 0be55239f..07a709d94 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -217,6 +217,51 @@ def test_build_async_chunk_outputs_slices_padded_axis_and_splits_channels() -> N assert client is None and torch.equal(inter_stage[0]["codes.audio"], req_codes[0]) +@pytest.mark.parametrize("async_chunk", [False, True]) +@pytest.mark.parametrize( + "hidden_padded,codes_padded", + [(False, False), (True, False), (False, True), (True, True)], + ids=["unpadded", "hidden_padded", "codes_padded", "both_padded"], +) +@pytest.mark.parametrize("lengths", [(1, 1, 1), (3, 1, 1)], ids=["decode", "mixed_prefill_decode"]) +def test_async_output_slices_request_payloads_with_graph_padding( + monkeypatch, mocker, async_chunk, hidden_padded, codes_padded, lengths +): + """Sync and async transfers must exclude other requests and graph padding.""" + from vllm.v1.worker.gpu.input_batch import InputBatch + + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + total = sum(lengths) + padded_total = 8 if hidden_padded or codes_padded else total + offsets = np.cumsum([0, *lengths]) + hidden_rows = padded_total if hidden_padded else total + code_rows = padded_total if codes_padded else total + hidden = torch.arange(hidden_rows * 4, dtype=torch.float32).reshape(hidden_rows, 4) + codes = torch.arange(code_rows * 16).reshape(code_rows, 16) + # Reference frames are request-local, even if their length matches the + # padded token count. They must not be sliced along the batch token axis. + refs = [torch.arange(padded_total * 16).reshape(padded_total, 16), torch.empty(0), torch.empty(0)] + batch = mocker.Mock(spec=InputBatch) + batch.query_start_loc_np = offsets + batch.num_scheduled_tokens = np.array(lengths) + batch.num_reqs = len(lengths) + batch.num_tokens_after_padding = padded_total + output = _async_output( + req_ids=[f"req-{i}" for i in range(len(lengths))], + sampler_output=SamplerOutput(torch.ones(len(lengths), 1, dtype=torch.long), None, None, None), + text_hidden=hidden, + multimodal_outputs={"codes": {"audio": codes, "ref": refs}}, + input_batch=batch, + async_chunk=async_chunk, + ).get_output() + + for i, payload in enumerate(output.inter_stage_outputs): + torch.testing.assert_close(payload["codes.audio"], codes[offsets[i] : offsets[i + 1]]) + torch.testing.assert_close(payload["codes.ref"], refs[i]) + if not async_chunk: + torch.testing.assert_close(payload["hidden"], hidden[offsets[i] : offsets[i + 1]]) + + def test_async_chunk_output_stages_mm_on_copy_stream_before_get_output(monkeypatch) -> None: monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) calls = [] diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 1af5c1f0a..8ff6e02f3 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -320,9 +320,12 @@ class OmniARModelRunner(OmniGPUModelRunner): query_start_loc_np: np.ndarray, num_scheduled_tokens: np.ndarray, num_reqs: int, + padded_total_tokens: int | None = None, ) -> list[dict[str, Any]]: - """Build pooler_output from already-CPU tensors.""" - total = hidden_cpu.shape[0] + """Build per-request payloads, excluding any CUDA Graph padding.""" + # Graph-padded hidden states and unpadded multimodal outputs can + # coexist. Use the scheduled token count to identify the real axis. + total = int(num_scheduled_tokens[:num_reqs].sum()) pooler: list[dict[str, Any]] = [] for i in range(num_reqs): start = int(query_start_loc_np[i]) @@ -335,6 +338,7 @@ class OmniARModelRunner(OmniGPUModelRunner): start=start, end=end, total_tokens=total, + padded_total_tokens=padded_total_tokens, ) pooler.append(flatten_payload(payload)) return pooler @@ -807,6 +811,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._query_start_loc_np, self._num_scheduled_tokens, self._num_reqs, + self._padded_total_tokens, ) pooler_payload = cast(list[dict[str, Any] | None], pooler_output) if pooler_output else None self.model_runner_output.pooler_output = pooler_payload