diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index f671e787d..f2c027dd0 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -14,11 +14,11 @@ RUN apt-get update && \ # Install vLLM wheel for the specific commit # URL format: https://wheels.vllm.ai// -# Current vLLM wheel commit: 6042e66cd5304fc043d96aaa0c22d56f939af320 +# Current vLLM wheel commit: a60985b07eaf0fefa4c353ad2309b745707591ed # Note: --prerelease=allow is required because nightly/commit wheels use # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default, # causing it to fall back to the older stable release from PyPI. -RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/6042e66cd5304fc043d96aaa0c22d56f939af320 +RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/a60985b07eaf0fefa4c353ad2309b745707591ed RUN uv pip uninstall --system nvidia-cublas-cu12 # Install vllm-omni into the same uv-managed Python environment used by the base image. diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 5eb25f169..1b1965355 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -86,7 +86,7 @@ def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.M data = {"large": "x" * 200} # Mock SHM return values - mock_handle = {"name": "test_shm", "size": 200} + mock_handle = {"name": "req_2", "size": 200} mock_write = mocker.MagicMock(return_value=mock_handle) monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write) diff --git a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml index 59bdc9df9..c11de3e5d 100644 --- a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml @@ -74,12 +74,13 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.15 + gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800 enforce_eager: true trust_remote_code: true enable_prefix_caching: false engine_output_type: audio - max_num_batched_tokens: 100000 + max_num_batched_tokens: 4096 + max_model_len: 4096 engine_input_source: [1] final_output: true final_output_type: audio diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 10538ddb4..7d6f27597 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -474,7 +474,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, - activate_lora: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -497,12 +496,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run - num_active_loras: Number of active LoRAs to capture for. - activate_lora: Backward-compatible override for LoRA activation. + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. """ - if activate_lora is None: - activate_lora = num_active_loras > 0 - mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: # The current dummy run only covers LM execution, so we can skip it. @@ -529,7 +525,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. - assert num_tokens <= self.scheduler_config.max_num_batched_tokens + assert num_tokens <= self.max_num_tokens max_num_reqs = self.scheduler_config.max_num_seqs if create_mixed_batch: assert not uniform_decode @@ -579,8 +575,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # `force_has_lora` is used for cudagraph capture; because LoRA is # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture - force_has_lora=activate_lora, - # Capture shape specialization for specific active LoRA counts. + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts force_num_active_loras=num_active_loras, ) ) @@ -653,8 +650,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): self.lora_config, num_scheduled_tokens, num_sampled_tokens, - activate_lora, remove_lora, + num_active_loras, ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens @@ -727,8 +724,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - if self.speculative_config and self.speculative_config.use_eagle(): + if self.speculative_config and ( + self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() + ): assert isinstance(self.drafter, EagleProposer) + assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. # Therefore only use cudagraphs if the main model uses PIECEWISE # NOTE(lucas): this is a hack, need to clean up. @@ -741,7 +741,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # lora cases when cudagraph_specialize_lora is enabled. This is a # short term mitigation for issue mentioned in # https://github.com/vllm-project/vllm/issues/28334 - if self.compilation_config.cudagraph_specialize_lora and activate_lora: + if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0: use_cudagraphs = False self.drafter.dummy_run( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index b8874133c..b16e21404 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -274,6 +274,11 @@ class OmniGPUModelRunner(GPUModelRunner): # Add new requests to the cached states. for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id + if req_id in self.requests: + req_state = self._update_streaming_request(req_id, new_req_data) + reqs_to_add.append(req_state) + continue + sampling_params = new_req_data.sampling_params pooling_params = new_req_data.pooling_params @@ -377,18 +382,15 @@ class OmniGPUModelRunner(GPUModelRunner): req_state.num_computed_tokens = num_computed_tokens if not is_last_rank: - # When using PP, the scheduler sends the sampled tokens back, - # because there's no direct communication between the first- - # stage worker and the last-stage worker. - new_token_ids = req_data.new_token_ids[i] - # Add the sampled token(s) from the previous step (if any). - # This doesn't include "unverified" tokens like spec tokens. - num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens - if num_new_tokens == 1: - # Avoid slicing list in most common case. - req_state.output_token_ids.append(new_token_ids[-1]) - elif num_new_tokens > 0: - req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) + if not req_data.new_token_ids: + new_token_ids: list[int] = [] + else: + new_token_ids = req_data.new_token_ids[i] + num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens + if num_new_tokens == 1: + req_state.output_token_ids.append(new_token_ids[-1]) + elif num_new_tokens > 0: + req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) elif num_output_tokens < len(req_state.output_token_ids): # Some output tokens were discarded due to a sync-KV-load # failure. Align the cached state. @@ -489,7 +491,6 @@ class OmniGPUModelRunner(GPUModelRunner): remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, - activate_lora: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -512,12 +513,9 @@ class OmniGPUModelRunner(GPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run - num_active_loras: Number of active LoRAs to capture for. - activate_lora: Backward-compatible override for LoRA activation. + num_active_loras: Number of distinct active LoRAs to capture for. + LoRA is activated when num_active_loras > 0. """ - if activate_lora is None: - activate_lora = num_active_loras > 0 - mm_config = self.vllm_config.model_config.multimodal_config if mm_config and mm_config.mm_encoder_only: # The current dummy run only covers LM execution, so we can skip it. @@ -544,7 +542,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Set num_scheduled_tokens based on num_tokens and max_num_seqs # for dummy run with LoRA so that the num_reqs collectively # has num_tokens in total. - assert num_tokens <= self.scheduler_config.max_num_batched_tokens + assert num_tokens <= self.max_num_tokens max_num_reqs = self.scheduler_config.max_num_seqs if create_mixed_batch: assert not uniform_decode @@ -594,8 +592,9 @@ class OmniGPUModelRunner(GPUModelRunner): # `force_has_lora` is used for cudagraph capture; because LoRA is # activated later in the context manager, but we need to know the # LoRA state when determining the batch descriptor for capture - force_has_lora=activate_lora, - # Capture shape specialization for specific active LoRA counts. + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts force_num_active_loras=num_active_loras, ) ) @@ -667,8 +666,8 @@ class OmniGPUModelRunner(GPUModelRunner): self.lora_config, num_scheduled_tokens, num_sampled_tokens, - activate_lora, remove_lora, + num_active_loras, ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens @@ -757,8 +756,11 @@ class OmniGPUModelRunner(GPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - if self.speculative_config and self.speculative_config.use_eagle(): + if self.speculative_config and ( + self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() + ): assert isinstance(self.drafter, EagleProposer) + assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. # Therefore only use cudagraphs if the main model uses PIECEWISE # NOTE(lucas): this is a hack, need to clean up. @@ -771,7 +773,7 @@ class OmniGPUModelRunner(GPUModelRunner): # lora cases when cudagraph_specialize_lora is enabled. This is a # short term mitigation for issue mentioned in # https://github.com/vllm-project/vllm/issues/28334 - if self.compilation_config.cudagraph_specialize_lora and activate_lora: + if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0: use_cudagraphs = False self.drafter.dummy_run(