anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Update vLLM wheel version and adjust GPU memory utilization in CI configuration

Browse Source 
- Updated the vLLM wheel commit hash in Dockerfile.ci to a60985b07eaf0fefa4c353ad2309b745707591ed.
- Modified the mock shared memory handle name in test_basic_connectors.py for clarity.
- Increased GPU memory utilization to 0.9 and adjusted max_num_batched_tokens to 4096 in qwen2_5_omni_ci.yaml to enable testing on H800.
- Refactored GPU model runner classes to remove deprecated activate_lora parameter and ensure compatibility with num_active_loras.
- Enhanced request handling in the OmniGPUModelRunner class to improve streaming request updates.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-03-02 03:08:45 +00:00
1 parent 26744bcb35
commit 473d3ecd29
5 files changed
+44 -41
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docker/Dockerfile.ci
+2 -2
View File
@@ -14,11 +14,11 @@ RUN apt-get update && \
# Install vLLM wheel for the specific commit
# URL format: https://wheels.vllm.ai/<commit_hash>/
# 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.
tests/distributed/omni_connectors/test_basic_connectors.py
+1 -1
View File
@@ -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)
tests/e2e/stage_configs/qwen2_5_omni_ci.yaml
+3 -2
View File
@@ -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
vllm_omni/worker/gpu_generation_model_runner.py
+12 -12
View File
@@ -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(
vllm_omni/worker/gpu_model_runner.py
+26 -24
View File
@@ -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(