From 43e05c6e13b5cda24875bdff80286aa4a16da98d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 3 Oct 2026 22:03:13 +0800 Subject: [PATCH] fix: align MRV2 profiling and capture with vLLM 0.31 Signed-off-by: tzhouam --- .../test_omni_generation_model_runner.py | 11 +++++ tests/worker_v2/test_omni_gpu_model_runner.py | 46 +++++++++++++++---- tests/worker_v2/test_omni_model_state.py | 10 ++++ .../model_states/omni_model_state.py | 3 ++ vllm_omni/worker_v2/omni_ar_model_runner.py | 2 + .../worker_v2/omni_generation_model_runner.py | 6 ++- vllm_omni/worker_v2/omni_model_runner.py | 15 ++++-- 7 files changed, 81 insertions(+), 12 deletions(-) diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py index 84927b3eb..df1d9f3ca 100644 --- a/tests/worker_v2/test_omni_generation_model_runner.py +++ b/tests/worker_v2/test_omni_generation_model_runner.py @@ -20,6 +20,17 @@ from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModel pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("randomize_inputs", [False, True]) +def test_profile_run_accepts_release_flag_without_running_codec(monkeypatch, randomize_inputs): + runner = object.__new__(OmniGenerationModelRunner) + runner.model = MagicMock(side_effect=AssertionError("profiling ran codec model")) + synchronize = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + runner.profile_run(randomize_inputs=randomize_inputs) + synchronize.assert_called_once_with() + runner.model.assert_not_called() + + class _FakeStagedField: def __init__(self, data): self.np = data diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py index dd4c760c6..a57ec8291 100644 --- a/tests/worker_v2/test_omni_gpu_model_runner.py +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -95,7 +95,8 @@ def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker): @pytest.mark.parametrize("output_form", ["tuple", "omni"]) -def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): +@pytest.mark.parametrize("profile_only", [False, True]) +def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form, profile_only): runner = object.__new__(OmniGPUModelRunner) hidden = torch.ones(1, 2) @@ -117,16 +118,21 @@ def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock()) runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc") - def assert_unwrapped(_self): + def assert_unwrapped(_self, *, profile_only: bool = False): assert torch.equal(runner.model.forward(), hidden) # forward unwrapped during capture + assert profile_only is expected_profile_only return 3 + expected_profile_only = profile_only with patch.object(GPUModelRunner, "capture_model", assert_unwrapped): - assert runner.capture_model() == 3 + assert runner.capture_model(profile_only=profile_only) == 3 assert runner.model.forward is original_forward # restored after capture assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]} - runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor) + if profile_only: + runner.model_state.capture_mtp_graphs.assert_not_called() + else: + runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor) @pytest.mark.parametrize( @@ -209,7 +215,7 @@ def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves(): runner.model_state = SimpleNamespace() captured = {} - def capture(_self): + def capture(_self, *, profile_only=False): # The graph manager stores (hidden, leaves) in its aux buffers. assert runner.use_aux_hidden_state_outputs is True hidden, leaves = runner.model.forward(4) @@ -262,7 +268,9 @@ def test_full_graph_output_without_contract_is_hidden_only(): @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"]) -def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): +@pytest.mark.parametrize("randomize_inputs", [False, True]) +@pytest.mark.parametrize("is_profile", [False, True]) +def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inputs, is_profile, monkeypatch): from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_runner import ExecuteModelState @@ -288,6 +296,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): runner.model_state.prepare_inputs.return_value = {} runner._omni_data_plane = object() runner.supports_mm_inputs = False + runner.vocab_size = 32 runner.lora_config = None runner.is_encoder_decoder = False runner.eplb = MagicMock() @@ -308,7 +317,15 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): cg_mode=CUDAGraphMode.NONE, num_reqs=1, num_tokens=1, num_active_loras=0, max_query_len=1 ) runner._dispatch_batch_descriptor = MagicMock(return_value=(batch_desc, None)) - monkeypatch.setattr(InputBatch, "make_dummy", lambda *args, **kwargs: input_batch) + make_dummy = MagicMock(return_value=input_batch) + monkeypatch.setattr(InputBatch, "make_dummy", make_dummy) + randomized = [] + + def randomize(tensor, low, high): + randomized.append((low, high)) + return tensor.fill_(7) + + monkeypatch.setattr(torch.Tensor, "random_", randomize) monkeypatch.setattr("vllm_omni.worker_v2.omni_model_runner.build_slot_mappings_by_layer", lambda *args: {}) for module in ("omni_model_runner", "omni_generation_model_runner"): monkeypatch.setattr(f"vllm_omni.worker_v2.{module}.set_forward_context", lambda *args, **kwargs: nullcontext()) @@ -318,7 +335,20 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): # Upstream _dummy_run always supplies valid_dummy_state_slots, and the # result must use the real upstream constructor rather than a mocked state. - assert runner.execute_model(scheduled, dummy_run=True, valid_dummy_state_slots=True) is None + assert ( + runner.execute_model( + scheduled, + dummy_run=True, + valid_dummy_state_slots=True, + randomize_inputs=randomize_inputs, + is_profile=is_profile, + ) + is None + ) + assert make_dummy.call_args.kwargs["is_padding"] is not is_profile + should_randomize = randomize_inputs and runner_kind != "generation" + assert randomized == ([(0, 32)] if should_randomize else []) + assert input_batch.input_ids.tolist() == ([7] if should_randomize else [1]) assert isinstance(runner.execute_model_state, ExecuteModelState) assert runner.execute_model_state.input_batch is input_batch assert runner.execute_model_state.cudagraph_stats is None diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 86764b691..1aa079d6e 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -22,6 +22,16 @@ from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _m pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_prepare_attn_forwards_release_model_specific_metadata(): + state = object.__new__(OmniModelState) + metadata = object() + output = object() + batch = SimpleNamespace() + with patch.object(DefaultModelState, "prepare_attn", return_value=output) as prepare: + assert state.prepare_attn(batch, None, (), None, [], None, model_specific_attn_metadata=metadata) is output + assert prepare.call_args.kwargs["model_specific_attn_metadata"] is metadata + + class _DummyInputBatch: is_prefilling_np: np.ndarray diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 34c3c1cdf..3d402b8ac 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -31,6 +31,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup @@ -447,6 +448,7 @@ class OmniModelState(DefaultModelState): kv_cache_config: KVCacheConfig, for_capture: bool = False, ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, ) -> dict[str, Any]: if ( for_capture @@ -471,6 +473,7 @@ class OmniModelState(DefaultModelState): kv_cache_config, for_capture=for_capture, ubatch_idx=ubatch_idx, + model_specific_attn_metadata=model_specific_attn_metadata, ) def prepare_inputs_embeds( diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 24002cfa4..11b14e2d9 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -121,6 +121,7 @@ class OmniARModelRunner(OmniGPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> Any: if not dummy_run: self._handle_kv_transfer_pre(scheduler_output) @@ -132,6 +133,7 @@ class OmniARModelRunner(OmniGPUModelRunner): is_profile=is_profile, context_len=context_len, valid_dummy_state_slots=valid_dummy_state_slots, + randomize_inputs=randomize_inputs, ) # ------------------------------------------------------------------ diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py index 2de7314fd..abde9a7db 100644 --- a/vllm_omni/worker_v2/omni_generation_model_runner.py +++ b/vllm_omni/worker_v2/omni_generation_model_runner.py @@ -292,7 +292,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): # ------------------------------------------------------------------ @torch.inference_mode() - def profile_run(self) -> None: + def profile_run(self, randomize_inputs: bool = False) -> None: """Generation models have no KV cache — skip profiling. Code2Wav shares GPU memory with the Talker stage (same device); @@ -317,6 +317,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: self._prepare_native_data_plane(scheduler_output) @@ -388,7 +389,10 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + is_padding=not is_profile, ) + # Codec IDs have model-specific bounds; keep valid dummy tokens + # even when upstream requests randomized language-model inputs. attn_metadata = None slot_mappings_by_layer = None diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 110f80a90..c07c47f89 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -391,7 +391,7 @@ class OmniGPUModelRunner(GPUModelRunner): } logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.") - def capture_model(self) -> int: + def capture_model(self, *, profile_only: bool = False) -> int: """Handle CUDA graph capture for Omni models. Tuple-returning models use PIECEWISE graphs; FULL replay requires @@ -424,12 +424,17 @@ class OmniGPUModelRunner(GPUModelRunner): use_aux = self.use_aux_hidden_state_outputs self.use_aux_hidden_state_outputs = use_aux or aux_outputs try: - result = super().capture_model() + result = super().capture_model(profile_only=profile_only) finally: self.model.forward = original_forward # type: ignore[assignment] self.use_aux_hidden_state_outputs = use_aux else: - result = super().capture_model() + result = super().capture_model(profile_only=profile_only) + + if profile_only: + # Upstream discards captures in its temporary profiling pool. + # Model-owned graphs must be recorded during the real capture. + return result capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None) if callable(capture_mtp): @@ -494,6 +499,7 @@ class OmniGPUModelRunner(GPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> Any: if not dummy_run: self._prepare_native_data_plane(scheduler_output) @@ -570,7 +576,10 @@ class OmniGPUModelRunner(GPUModelRunner): batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + is_padding=not is_profile, ) + if randomize_inputs: + input_batch.input_ids.random_(0, self.vocab_size) if not skip_attn_for_dummy_run: block_tables, slot_mappings = self.prepare_dummy_attn(input_batch, valid_dummy_state_slots) if context_len: