From 11a5a914cab71fdde5a935828161acdabe7a3e15 Mon Sep 17 00:00:00 2001 From: XU Mingshi <91017482+mxuax@users.noreply.github.com> Date: Fri, 6 Feb 2026 17:28:49 +0800 Subject: [PATCH 1/9] [Bugfix] Fix wan2.2 ti2v (#1221) Signed-off-by: mxuax Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../diffusion/distributed/sp_sharding.py | 5 ++ .../models/wan2_2/wan2_2_transformer.py | 85 ++++++++++++++++--- 2 files changed, 76 insertions(+), 14 deletions(-) diff --git a/vllm_omni/diffusion/distributed/sp_sharding.py b/vllm_omni/diffusion/distributed/sp_sharding.py index 12520183c..69287d25b 100644 --- a/vllm_omni/diffusion/distributed/sp_sharding.py +++ b/vllm_omni/diffusion/distributed/sp_sharding.py @@ -63,6 +63,11 @@ def sp_shard( f"world_size ({world_size}) for sequence parallel sharding." ) + if size < world_size: + raise ValueError( + f"Tensor size along dim {dim} ({size}) must be >= world_size ({world_size}). Tensor shape: {tensor.shape}" + ) + return tensor.chunk(world_size, dim=dim)[rank] diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 328a64024..4de44119f 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -296,6 +296,47 @@ class WanTimeTextImageEmbedding(nn.Module): return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image +class TimestepProjPrepare(nn.Module): + """Prepares timestep_proj for sequence parallel in TI2V models. + + Encapsulates the unflatten operation for timestep_proj to enable _sp_plan sharding. + """ + + def forward( + self, + timestep_proj: torch.Tensor, + ts_seq_len: int | None, + ) -> torch.Tensor: + if ts_seq_len is not None: + # TI2V mode: [batch, seq_len, 6, inner_dim] + timestep_proj = timestep_proj.unflatten(2, (6, -1)) + else: + # T2V mode: [batch, 6, inner_dim] + timestep_proj = timestep_proj.unflatten(1, (6, -1)) + return timestep_proj + + +class OutputScaleShiftPrepare(nn.Module): + """Prepares output scale/shift for SP sharding in TI2V models.""" + + def __init__(self, inner_dim: int): + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if temb.ndim == 3: + # TI2V: [B, seq, D] -> 3D outputs for SP sharding + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift = shift.squeeze(2) + scale = scale.squeeze(2) + else: + # T2V: [B, D] -> 2D outputs (skip SP sharding via expected_dims=3) + shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) + shift = shift.squeeze(1) + scale = scale.squeeze(1) + return shift, scale + + class WanSelfAttention(nn.Module): """ Optimized self-attention module using vLLM layers. @@ -679,6 +720,7 @@ class WanTransformer3DModel(nn.Module): # # The _sp_plan specifies sharding/gathering at module boundaries: # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True + # - timestep_proj_prepare: Split timestep_proj for TI2V models (4D tensor) # - blocks.0: Split hidden_states input at the first transformer block # - proj_out: Gather outputs after the final projection layer # @@ -689,11 +731,22 @@ class WanTransformer3DModel(nn.Module): 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # freqs_cos [1, seq, 1, dim] 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # freqs_sin [1, seq, 1, dim] }, + # Shard timestep_proj for TI2V models (4D tensor: [batch, seq_len, 6, inner_dim]) + # This is only active when ts_seq_len is not None (TI2V mode) + # Output is a single tensor, shard along dim=1 (sequence dimension) + "timestep_proj_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # [B, seq, 6, dim] + }, # Shard hidden_states at first transformer block input # (after patch_embedding + flatten + transpose) "blocks.0": { "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), # [B, seq, dim] }, + # Shard output scale/shift for TI2V (3D); T2V outputs 2D and skips sharding + "output_scale_shift_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + }, # Gather at proj_out (final linear projection before unpatchify) "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } @@ -774,7 +827,10 @@ class WanTransformer3DModel(nn.Module): # 4. Output norm & projection self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) - self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + # SP helper modules + self.timestep_proj_prepare = TimestepProjPrepare() + self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim) @property def dtype(self) -> torch.dtype: @@ -814,10 +870,10 @@ class WanTransformer3DModel(nn.Module): temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len ) - if ts_seq_len is not None: - timestep_proj = timestep_proj.unflatten(2, (6, -1)) - else: - timestep_proj = timestep_proj.unflatten(1, (6, -1)) + # Prepare timestep_proj via TimestepProjPrepare module + # _sp_plan will shard timestep_proj via split_output=True (when ts_seq_len is not None) + # This ensures timestep_proj sequence dimension matches sharded hidden_states + timestep_proj = self.timestep_proj_prepare(timestep_proj, ts_seq_len) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) @@ -827,15 +883,12 @@ class WanTransformer3DModel(nn.Module): hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) # Output norm, projection & unpatchify - if temb.ndim == 3: - shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) - shift = shift.squeeze(2) - scale = scale.squeeze(2) - else: - shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) - + shift, scale = self.output_scale_shift_prepare(temb) shift = shift.to(hidden_states.device) scale = scale.to(hidden_states.device) + if shift.ndim == 2: # T2V mode: unsqueeze for broadcasting + shift = shift.unsqueeze(1) + scale = scale.unsqueeze(1) hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) hidden_states = self.proj_out(hidden_states) @@ -867,8 +920,6 @@ class WanTransformer3DModel(nn.Module): tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() # Stacked params mapping for self-attention QKV fusion - # Format: (param_name, shard_name, shard_id) - # Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention) stacked_params_mapping = [ # self-attention QKV fusion (".attn1.to_qkv", ".attn1.to_q", "q"), @@ -876,10 +927,16 @@ class WanTransformer3DModel(nn.Module): (".attn1.to_qkv", ".attn1.to_v", "v"), ] + # Remap scale_shift_table to new module location + weight_name_remapping = { + "scale_shift_table": "output_scale_shift_prepare.scale_shift_table", + } + params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: + name = weight_name_remapping.get(name, name) original_name = name lookup_name = name From 10f841bb9840dae3cb70b68bcdace7cadd52513b Mon Sep 17 00:00:00 2001 From: ApsarasX Date: Fri, 6 Feb 2026 21:44:51 +0800 Subject: [PATCH 2/9] [Bugfix] Fix '--max-generated-image-size' cli args type (#1249) Signed-off-by: ApsarasX --- vllm_omni/entrypoints/cli/serve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index acefc081b..ac03de143 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -245,7 +245,7 @@ class OmniServeCommand(CLISubcommand): # Diffusion model mixed precision omni_config_group.add_argument( "--max-generated-image-size", - type=float, + type=int, help="The max size of generate image (height * width).", ) return serve_parser From 87847a214ab1bd7e259e196df53812d6a18cb565 Mon Sep 17 00:00:00 2001 From: ApsarasX Date: Sat, 7 Feb 2026 07:09:54 +0800 Subject: [PATCH 3/9] [Bugfix] Ensure seed=0 is correctly handled in image edit (#1248) Signed-off-by: ApsarasX --- .../openai_api/test_image_server.py | 71 +++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 6 +- 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 3df484b1e..6ff6782ae 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -338,6 +338,25 @@ def test_with_seed(test_client): assert response.status_code == 200 +def test_with_seed_zero(test_client): + """Test with seed=0 for reproducibility""" + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a tree", + "seed": 0, + "size": "1024x1024", + }, + ) + assert response.status_code == 200 + engine = test_client.app.state.engine_client + captured = engine.captured_sampling_params_list[0] + # Verify that seed=0 is correctly passed + assert captured.seed == 0, ( + f"Expected seed=0, but got seed={captured.seed}. This indicates the bug where seed=0 is treated as falsy." + ) + + def test_with_custom_parameters(test_client): """Test with custom diffusion parameters""" response = test_client.post( @@ -814,3 +833,55 @@ def test_image_edit_compression_png(async_omni_test_client): assert len(img_bytes_10) < len(img_bytes_50) assert len(img_bytes_50) < len(img_bytes_100) + + +def test_image_edit_with_seed_zero(async_omni_test_client): + """Test that seed=0 is correctly handled in image editing. + + Previously, seed=0 was incorrectly replaced by a random seed due to the + falsy value check using `or` operator. This test ensures seed=0 is + properly passed through to the sampling parameters in image editing. + """ + img_bytes_1 = make_test_image_bytes((16, 16)) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes_1)], + data={ + "prompt": "edit this image", + "seed": 0, + }, + ) + assert response.status_code == 200 + engine = async_omni_test_client.app.state.engine_client + captured_sampling_params = engine.captured_sampling_params_list[-1] + # Verify that seed=0 is correctly passed + assert captured_sampling_params.seed == 0, ( + f"Expected seed=0, but got seed={captured_sampling_params.seed}. " + "This indicates the bug where seed=0 is treated as falsy." + ) + + +def test_image_edit_with_seed_zero_single_stage(test_client): + """Test that seed=0 is correctly handled in image editing (single stage). + + Test seed=0 handling in image editing with single stage path. + """ + img_bytes_1 = make_test_image_bytes((16, 16)) + + response = test_client.post( + "/v1/images/edits", + files=[("image", img_bytes_1)], + data={ + "prompt": "edit this image", + "seed": 0, + }, + ) + assert response.status_code == 200 + engine = test_client.app.state.engine_client + captured_sampling_params = engine.captured_sampling_params_list[0] + # Verify that seed=0 is correctly passed + assert captured_sampling_params.seed == 0, ( + f"Expected seed=0, but got seed={captured_sampling_params.seed}. " + "This indicates the bug where seed=0 is treated as falsy." + ) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index fb52c7e46..b5e48ab95 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -968,7 +968,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) # a proper generator is initialized in the backend. # This fixes issues where using the default global generator # might produce blurry images in some environments. - _update_if_not_none(gen_params, "seed", random.randint(0, 2**32 - 1) if request.seed is None else request.seed) + _update_if_not_none( + gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1) + ) request_id = f"img_gen_{uuid.uuid4().hex}" @@ -1126,7 +1128,7 @@ async def edit_images( # a proper generator is initialized in the backend. # This fixes issues where using the default global generator # might produce blurry images in some environments. - _update_if_not_none(gen_params, "seed", seed or random.randint(0, 2**32 - 1)) + _update_if_not_none(gen_params, "seed", seed if seed is not None else random.randint(0, 2**32 - 1)) # 4. Generate images using AsyncOmni (multi-stage mode) request_id = f"img_edit_{int(time.time())}" From a43dbdc06480ce8106cb986a65970f650e46fdfe Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 7 Feb 2026 21:14:53 +0800 Subject: [PATCH 4/9] [Docs] Add example image download step to Image-To-Video examples (#1258) Signed-off-by: lishunyang --- .../examples/offline_inference/image_to_video.md | 16 ++++++++++++---- .../offline_inference/image_to_video/README.md | 16 ++++++++++++---- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index 60f00a95b..c83d39435 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -7,12 +7,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image- ## Local CLI Usage +Download the example image: + +```bash +wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +``` + ### Wan2.2-I2V-A14B-Diffusers (MoE) + ```bash python image_to_video.py \ --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ - --image input.png \ - --prompt "A cat playing with yarn, smooth motion" \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ --negative_prompt "" \ --height 480 \ --width 832 \ @@ -27,11 +34,12 @@ python image_to_video.py \ ``` ### Wan2.2-TI2V-5B-Diffusers (Unified) + ```bash python image_to_video.py \ --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --image input.png \ - --prompt "A cat playing with yarn, smooth motion" \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ --negative_prompt "" \ --height 480 \ --width 832 \ diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index a1355dab6..2d8b33d0d 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -4,12 +4,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image- ## Local CLI Usage +Download the example image: + +```bash +wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +``` + ### Wan2.2-I2V-A14B-Diffusers (MoE) + ```bash python image_to_video.py \ --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ - --image input.png \ - --prompt "A cat playing with yarn, smooth motion" \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ --negative_prompt "" \ --height 480 \ --width 832 \ @@ -24,11 +31,12 @@ python image_to_video.py \ ``` ### Wan2.2-TI2V-5B-Diffusers (Unified) + ```bash python image_to_video.py \ --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --image input.png \ - --prompt "A cat playing with yarn, smooth motion" \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ --negative_prompt "" \ --height 480 \ --width 832 \ From 53a44c56fe07055dc049663c1495856ad9e46177 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Sat, 7 Feb 2026 11:35:49 -0500 Subject: [PATCH 5/9] [Bugfix] Fix padding bug in 12Hz tokenizer ConvTranspose1d decode (#1241) Signed-off-by: linyueqian --- .../tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 31abe254b..78fc965fe 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -197,12 +197,13 @@ class Qwen3TTSTokenizerV2CausalTransConvNet(nn.Module): self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride) pad = kernel_size - stride - self.left_pad = math.ceil(pad) - self.right_pad = pad = self.left_pad + self.left_pad = 0 + self.right_pad = int(pad) def forward(self, hidden_state): hidden_state = self.conv(hidden_state) - hidden_state = hidden_state[..., self.left_pad : hidden_state.shape[-1] - self.right_pad] + if self.right_pad > 0: + hidden_state = hidden_state[..., : hidden_state.shape[-1] - self.right_pad] return hidden_state.contiguous() From accf33411251c51e3e391cb6eaa67b2f91faee6a Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Sat, 7 Feb 2026 11:39:09 -0500 Subject: [PATCH 6/9] [bugfix] Fix multimodal_output property to check completion outputs where audio data is attached (#1203) Signed-off-by: linyueqian --- .../online_serving/qwen3_tts/openai_speech_client.py | 11 ++++++++--- vllm_omni/outputs.py | 8 +++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/examples/online_serving/qwen3_tts/openai_speech_client.py b/examples/online_serving/qwen3_tts/openai_speech_client.py index 4d7046094..ea45a7482 100644 --- a/examples/online_serving/qwen3_tts/openai_speech_client.py +++ b/examples/online_serving/qwen3_tts/openai_speech_client.py @@ -111,9 +111,14 @@ def run_tts_generation(args) -> None: print(response.text) return - if response.content.decode("utf-8").startswith('{"error"'): - print(f"Error: {response.content.decode('utf-8')}") - return + # Check for JSON error response (only if content is valid UTF-8 text) + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass # Binary audio data, not an error # Save audio response output_path = args.output or "tts_output.wav" diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 6cba9d96a..64569846a 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -124,10 +124,16 @@ class OmniRequestOutput: def multimodal_output(self) -> dict[str, Any]: """Return multimodal output from the underlying request output or local field. - For pipeline outputs, this proxies to request_output.multimodal_output. + For pipeline outputs, this checks completion outputs first, then request_output. For diffusion outputs, this returns the local _multimodal_output field. """ if self.request_output is not None: + # Check completion outputs first (where multimodal_output is attached) + if self.request_output.outputs: + for output in self.request_output.outputs: + mm = getattr(output, "multimodal_output", None) + if mm: + return mm return getattr(self.request_output, "multimodal_output", {}) return self._multimodal_output From 2f97db17a3bb6ca51c2b20030f5b88456803a46a Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sun, 8 Feb 2026 07:10:49 +0800 Subject: [PATCH 7/9] [Doc] Update QA relevant to quantization (#1257) Signed-off-by: lishunyang --- docs/usage/faq.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/usage/faq.md b/docs/usage/faq.md index e109439e8..c080eae40 100644 --- a/docs/usage/faq.md +++ b/docs/usage/faq.md @@ -22,7 +22,7 @@ A: At first, you can check current [issues](https://github.com/vllm-project/vllm > Q: Does vLLM-Omni support AWQ or any other quantization? -A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution! +A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057). > Q: Does vLLM-Omni support multimodal streaming input and output? From 01e0a8f284576cff5c7962a2261dfa1036b6b57e Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sun, 8 Feb 2026 20:36:48 +0800 Subject: [PATCH 8/9] [Bugfix] Fix Doc link Rrror (#1263) Signed-off-by: lishunyang --- docs/user_guide/examples/offline_inference/text_to_image.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index c0dbfd141..2e8739949 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -51,7 +51,8 @@ if __name__ == "__main__": For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support - batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts. + batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md + ) to let the model accept a longer batch of prompts. Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM. This is useful for models that support negative prompts. From 59024bee62df5ff8660d59fe27f87445f36d4d41 Mon Sep 17 00:00:00 2001 From: dsinghvi Date: Sun, 8 Feb 2026 20:13:02 +0530 Subject: [PATCH 9/9] Process-Scoped GPU Memory Accounting (#1204) Signed-off-by: Divyansh Singhvi --- tests/worker/test_process_gpu_memory.py | 154 +++++++ vllm_omni/entrypoints/omni_stage.py | 513 ++++++++-------------- vllm_omni/worker/base.py | 102 +++++ vllm_omni/worker/gpu_ar_worker.py | 4 +- vllm_omni/worker/gpu_generation_worker.py | 4 +- vllm_omni/worker/gpu_memory_utils.py | 124 ++++++ 6 files changed, 563 insertions(+), 338 deletions(-) create mode 100644 tests/worker/test_process_gpu_memory.py create mode 100644 vllm_omni/worker/base.py create mode 100644 vllm_omni/worker/gpu_memory_utils.py diff --git a/tests/worker/test_process_gpu_memory.py b/tests/worker/test_process_gpu_memory.py new file mode 100644 index 000000000..d8e0b969d --- /dev/null +++ b/tests/worker/test_process_gpu_memory.py @@ -0,0 +1,154 @@ +"""Tests for process-scoped GPU memory accounting.""" + +import os +from unittest import mock + +import pytest + + +class TestParseCudaVisibleDevices: + def test_empty(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {}, clear=True): + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + assert parse_cuda_visible_devices() == [] + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + assert parse_cuda_visible_devices() == [] + + def test_integer_indices(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}): + assert parse_cuda_visible_devices() == [2, 3, 5] + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}): + assert parse_cuda_visible_devices() == [0] + + def test_uuids(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + uuid1 = "GPU-12345678-1234-1234-1234-123456789abc" + uuid2 = "GPU-87654321-4321-4321-4321-cba987654321" + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}): + assert parse_cuda_visible_devices() == [uuid1, uuid2] + + def test_mig_ids(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0" + mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0" + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}): + assert parse_cuda_visible_devices() == [mig1, mig2] + + def test_spaces(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}): + assert parse_cuda_visible_devices() == [2, 3, 5] + + +class TestGetProcessGpuMemory: + @pytest.mark.skipif(not os.path.exists("/dev/nvidia0"), reason="No GPU") + def test_returns_memory_for_current_process(self): + import torch + + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + device = torch.device("cuda:0") + tensor = torch.zeros(1000, 1000, device=device) + + memory = get_process_gpu_memory(0) + assert memory >= 0 + + del tensor + torch.cuda.empty_cache() + + def test_raises_on_invalid_device(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1), + ): + with pytest.raises(RuntimeError, match="Invalid GPU device"): + get_process_gpu_memory(5) + + def test_returns_zero_when_process_not_found(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), + ): + memory = get_process_gpu_memory(0) + assert memory == 0 + + def test_uses_uuid_when_provided(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + uuid = "GPU-12345678-1234-1234-1234-123456789abc" + mock_handle = mock.MagicMock() + + with ( + mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid, + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), + ): + memory = get_process_gpu_memory(0) + assert memory == 0 + mock_by_uuid.assert_called_once_with(uuid) + + def test_raises_on_invalid_uuid(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + uuid = "GPU-invalid-uuid" + + with ( + mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")), + ): + with pytest.raises(RuntimeError, match="Failed to get NVML handle"): + get_process_gpu_memory(0) + + def test_returns_none_on_nvml_init_failure(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + ): + result = get_process_gpu_memory(0) + assert result is None + + +class TestIsProcessScopedMemoryAvailable: + def test_returns_true_when_nvml_works(self): + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + ): + assert is_process_scoped_memory_available() is True + + def test_returns_false_when_nvml_fails(self): + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + with mock.patch( + "vllm_omni.worker.gpu_memory_utils.nvmlInit", + side_effect=Exception("NVML unavailable"), + ): + assert is_process_scoped_memory_available() is False diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index a2070f89d..4935c43cc 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -16,6 +16,7 @@ import sys import time import traceback from collections.abc import Sequence +from contextlib import contextmanager from dataclasses import fields from typing import Any, Literal, cast @@ -55,6 +56,144 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +@contextmanager +def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = 300): + """Acquire device locks for sequential init if NVML is unavailable. + + If process-scoped memory tracking is available (NVML works), stages can + safely initialize concurrently — each measures only its own GPU memory. + Otherwise, fall back to file-based locks to serialize initialization. + """ + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + if is_process_scoped_memory_available(): + logger.debug("NVML process-scoped memory available — concurrent init is safe, skipping locks") + yield + return + + logger.debug("NVML unavailable — using sequential init locks") + + from vllm_omni.platforms import current_omni_platform + + # Get all parallel sizes from engine_args or parallel_config (defaults to 1) + if "parallel_config" in engine_args: + parallel_config = engine_args["parallel_config"] + tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) + pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) + data_parallel_size = parallel_config.get("data_parallel_size", 1) + prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) + sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) + cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + else: + tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) + pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) + data_parallel_size = engine_args.get("data_parallel_size", 1) + prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) + sequence_parallel_size = 1 + cfg_parallel_size = 1 + + num_devices_per_stage = ( + tensor_parallel_size + * pipeline_parallel_size + * data_parallel_size + * prefill_context_parallel_size + * sequence_parallel_size + * cfg_parallel_size + ) + + # Get physical device IDs from device control env var + device_control_env = current_omni_platform.device_control_env_var + visible_devices_str = os.environ.get(device_control_env) + physical_devices = [] + + if visible_devices_str: + try: + physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] + except (ValueError, IndexError): + pass + + if not physical_devices: + num_devices = current_omni_platform.get_device_count() + physical_devices = list(range(num_devices)) + + num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) + devices_to_lock = sorted(physical_devices[:num_devices_to_lock]) + + logger.debug( + "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", + tensor_parallel_size, + pipeline_parallel_size, + data_parallel_size, + prefill_context_parallel_size, + sequence_parallel_size, + cfg_parallel_size, + num_devices_to_lock, + devices_to_lock, + ) + + # Acquire exclusive locks for all devices using fcntl.flock + wait_start = time.time() + acquired_lock_fds = [] + + for device_id in devices_to_lock: + lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" + lock_acquired = False + + while not lock_acquired: + try: + lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644) + + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + os.ftruncate(lock_fd, 0) + os.write(lock_fd, f"{os.getpid()}\n".encode()) + os.fsync(lock_fd) + lock_acquired = True + acquired_lock_fds.append(lock_fd) + logger.debug("Acquired exclusive lock for device %s", device_id) + except BlockingIOError: + os.close(lock_fd) + + if time.time() - wait_start > stage_init_timeout: + logger.warning( + "Timeout waiting for device %s initialization lock, proceeding anyway", + device_id, + ) + break + + time.sleep(0.1) + except OSError as e: + logger.debug( + "Failed to acquire lock for device %s: %s, continuing anyway", + device_id, + e, + ) + try: + os.close(lock_fd) + except (OSError, NameError): + pass + break + + # Set FD_CLOEXEC to prevent child processes from inheriting locks + for lock_fd in acquired_lock_fds: + try: + flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) + fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + except (OSError, ValueError): + pass + + try: + yield + finally: + for lock_fd in acquired_lock_fds: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + logger.debug("Released initialization lock (fd=%s)", lock_fd) + except (OSError, ValueError): + pass + + def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: worker_type = engine_args.pop("worker_type", None) if not worker_type: @@ -541,183 +680,32 @@ def _stage_worker( except Exception as e: logger.warning("Device setup failed: %s", e) - # Sequential initialization on the same device to avoid memory calculation errors - # when multiple instances start simultaneously - # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage - lock_files = [] - try: - # Get all parallel sizes from engine_args or parallel_config (defaults to 1) - if "parallel_config" in engine_args: - parallel_config = engine_args["parallel_config"] - tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) - pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) - data_parallel_size = parallel_config.get("data_parallel_size", 1) - prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) - sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) - cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + # Use sequential init locks only when NVML is unavailable + with _sequential_init_lock(engine_args, stage_init_timeout): + # Init engine based on stage_type + logger.debug( + "[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()) + ) + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id + if stage_type == "diffusion": + engine_args.pop("model_stage", None) + engine_args.pop("model", None) + stage_engine = OmniDiffusion( + model=model, + stage_id=stage_id, + engine_input_source=stage_payload.get("engine_input_source", []), + **engine_args, + ) else: - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) - sequence_parallel_size = 1 # not use in omni model - cfg_parallel_size = 1 # not used in omni model - - # Calculate total number of devices needed for this stage - # For a single stage worker: - # - TP: splits model across devices (always needed) - # - PP: splits layers across pipeline stages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # - SP: sequence parallelism, typically uses TP devices - # - CFG: Classifier-Free Guidance parallelism for diffusion models - # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size - * pipeline_parallel_size - * data_parallel_size - * prefill_context_parallel_size - * sequence_parallel_size - * cfg_parallel_size - ) - - # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) - # After set_stage_devices, this env var is set to physical device(s) - device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = _os.environ.get(device_control_env) - physical_devices = [] - - if visible_devices_str: - try: - physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] - except (ValueError, IndexError): - pass - - if not physical_devices: - # Fallback: use logical device count if device control env var not set - num_devices = current_omni_platform.get_device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - - logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - sequence_parallel_size, - cfg_parallel_size, - num_devices_to_lock, - devices_to_lock, - ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, proceeding anyway", - device_id, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", - device_id, - e, - ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass - break - - lock_files = acquired_lock_fds - - # Set FD_CLOEXEC on all lock file descriptors to prevent child processes - # (e.g., EngineCore) from inheriting them, which would cause deadlock - for lock_fd in acquired_lock_fds: - try: - flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) - fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) - except (OSError, ValueError): - pass - except Exception as e: - logger.debug( - "[Stage-%s] Failed to set up sequential initialization lock: %s", - stage_id, - e, - ) - - # Init engine based on stage_type - logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) - if engine_args.get("async_chunk", False): - logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) - stage_connector_spec = {} - for v in connectors_config.values(): - stage_connector_spec = dict(v.get("spec", {})) - break - engine_args["stage_connector_spec"] = stage_connector_spec - engine_args["stage_id"] = stage_id - if stage_type == "diffusion": - engine_args.pop("model_stage", None) - engine_args.pop("model", None) - stage_engine = OmniDiffusion( - model=model, - stage_id=stage_id, - engine_input_source=stage_payload.get("engine_input_source", []), - **engine_args, - ) - else: - # Default to LLM engine - stage_engine = OmniLLM(model=model, **engine_args) - - # Release all locks AFTER engine initialization completes - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass - lock_files = [] # Clear after release + # Default to LLM engine + stage_engine = OmniLLM(model=model, **engine_args) logger.debug("Engine initialized") # Initialize OmniConnectors if configured @@ -1086,155 +1074,23 @@ async def _stage_worker_async( return connectors = built_connectors - # Sequential initialization on the same device to avoid memory calculation errors - # when multiple instances start simultaneously - # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage - lock_files = [] - try: - # Get all parallel sizes from engine_args or parallel_config (defaults to 1) - if "parallel_config" in engine_args: - parallel_config = engine_args["parallel_config"] - tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) - pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) - data_parallel_size = parallel_config.get("data_parallel_size", 1) - prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) - sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) - cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) - else: - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) - sequence_parallel_size = 1 # not use in omni model - cfg_parallel_size = 1 # not used in omni model - - # Calculate total number of devices needed for this stage - # For a single stage worker: - # - TP: splits model across devices (always needed) - # - PP: splits layers across pipeline stages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # - SP: sequence parallelism, typically uses TP devices - # - CFG: Classifier-Free Guidance parallelism for diffusion models - # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size - * pipeline_parallel_size - * data_parallel_size - * prefill_context_parallel_size - * sequence_parallel_size - * cfg_parallel_size - ) - - # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) - # After set_stage_devices, this env var is set to physical device(s) - device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = _os.environ.get(device_control_env) - physical_devices = [] - - if visible_devices_str: - try: - physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] - except (ValueError, IndexError): - pass - - if not physical_devices: - # Fallback: use logical device count if device control env var not set - num_devices = current_omni_platform.get_device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - + # Use sequential init locks only when NVML is unavailable + with _sequential_init_lock(engine_args, stage_init_timeout): + # Init engine based on stage_type logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - sequence_parallel_size, - cfg_parallel_size, - num_devices_to_lock, - devices_to_lock, + "[Stage-%s] Initializing %s engine with args keys=%s", + stage_id, + stage_type, + list(engine_args.keys()), ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s", - device_id, - stage_init_timeout, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", - device_id, - e, - ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass - break - - lock_files = acquired_lock_fds - except Exception as e: - logger.debug("Failed to set up sequential initialization lock: %s", e) - - # Init engine based on stage_type - logger.debug( - "[Stage-%s] Initializing %s engine with args keys=%s", - stage_id, - stage_type, - list(engine_args.keys()), - ) - if engine_args.get("async_chunk", False): - logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) - stage_connector_spec = {} - for v in connectors_config.values(): - stage_connector_spec = dict(v.get("spec", {})) - break - engine_args["stage_connector_spec"] = stage_connector_spec - engine_args["stage_id"] = stage_id - try: + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id if stage_type == "diffusion": # For diffusion, we need to extract diffusion-specific config od_config = _build_od_config(engine_args, model) @@ -1261,17 +1117,6 @@ async def _stage_worker_async( usage_context=usage_context, engine_args=omni_engine_args, ) - finally: - # Release all locks by closing file descriptors - # Locks are automatically released when file descriptors are closed - # or when process dies - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass omni_stage.set_async_engine(stage_engine) if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats: diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py new file mode 100644 index 000000000..2e7d07858 --- /dev/null +++ b/vllm_omni/worker/base.py @@ -0,0 +1,102 @@ +"""Base worker class for vLLM-Omni with process-scoped GPU memory accounting.""" + +from __future__ import annotations + +import os + +import torch +from vllm.logger import init_logger +from vllm.utils.mem_utils import format_gib, memory_profiling +from vllm.v1.worker.gpu_worker import Worker as GPUWorker + +from vllm_omni.worker.gpu_memory_utils import ( + get_process_gpu_memory, +) + +logger = init_logger(__name__) + + +class OmniGPUWorkerBase(GPUWorker): + """Base GPU worker for vLLM-Omni with process-scoped memory accounting. + + This class overrides determine_available_memory() to use per-process GPU + memory tracking via pynvml, allowing multiple stages to initialize + concurrently on the same GPU without memory accounting interference. + """ + + @torch.inference_mode() + def determine_available_memory(self) -> int: + """Process-scoped GPU memory profiling for concurrent stage initialization. + + Algorithm: + 1. requested_memory = total_gpu_memory * gpu_memory_utilization + (computed in init_device from cache_config) + + 2. process_memory = memory used by THIS process only (via pynvml) + - Uses nvmlDeviceGetComputeRunningProcesses to get per-PID memory + - Supports CUDA_VISIBLE_DEVICES with indices, UUIDs, or MIG IDs + + 3. available_kv_cache = requested_memory - process_memory + + Fallback: + If NVML is unavailable, falls back to profiling data: + available = requested - (weights + activations + non_torch) + """ + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: + self.model_runner.profile_run() + logger.info( + "Using explicit kv_cache_memory_bytes: %s GiB", + format_gib(kv_cache_memory_bytes), + ) + return kv_cache_memory_bytes + + with memory_profiling( + self.init_snapshot, + weights_memory=int(self.model_runner.model_memory_usage), + ) as profile_result: + self.model_runner.profile_run() + + self.non_torch_memory = profile_result.non_torch_increase + self.peak_activation_memory = profile_result.torch_peak_increase + + process_memory = get_process_gpu_memory(self.local_rank) + + if process_memory is not None: + # NVML available: use per-process memory + self.available_kv_cache_memory_bytes = max(0, self.requested_memory - process_memory) + logger.debug( + "Process-scoped memory (PID %d, GPU %d): requested=%s, used=%s, available=%s", + os.getpid(), + self.local_rank, + format_gib(self.requested_memory), + format_gib(process_memory), + format_gib(self.available_kv_cache_memory_bytes), + ) + logger.info_once( + "Available KV cache memory: %s GiB (process-scoped)", + format_gib(self.available_kv_cache_memory_bytes), + scope="local", + ) + else: + # NVML unavailable: use profiling data as conservative fallback + profiled_usage = ( + int(self.model_runner.model_memory_usage) + + profile_result.torch_peak_increase + + profile_result.non_torch_increase + ) + self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage) + logger.debug( + "Profiling fallback (PID %d, GPU %d): requested=%s, profiled=%s, available=%s", + os.getpid(), + self.local_rank, + format_gib(self.requested_memory), + format_gib(profiled_usage), + format_gib(self.available_kv_cache_memory_bytes), + ) + logger.info_once( + "Available KV cache memory: %s GiB (profiling fallback)", + format_gib(self.available_kv_cache_memory_bytes), + scope="local", + ) + + return int(self.available_kv_cache_memory_bytes) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 98b5afdda..30edba643 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -7,18 +7,18 @@ from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUARWorker(OmniWorkerMixin, GPUWorker): +class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU worker for autoregressive omni model stages. Extends the base GPUWorker to initialize and manage autoregressive diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 91f90d6b4..8ba4efa3d 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -7,18 +7,18 @@ from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUGenerationWorker(OmniWorkerMixin, GPUWorker): +class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU Worker for Generation model (non-autoregressive waveform generation). Usage in stage config: diff --git a/vllm_omni/worker/gpu_memory_utils.py b/vllm_omni/worker/gpu_memory_utils.py new file mode 100644 index 000000000..d005c5e58 --- /dev/null +++ b/vllm_omni/worker/gpu_memory_utils.py @@ -0,0 +1,124 @@ +"""NVML-based per-process GPU memory utilities. + +Shared across worker types (OmniGPUWorkerBase, DiffusionWorker, etc.) +for process-scoped GPU memory accounting. +""" + +from __future__ import annotations + +import os + +from vllm.logger import init_logger +from vllm.third_party.pynvml import ( + nvmlDeviceGetComputeRunningProcesses, + nvmlDeviceGetHandleByIndex, + nvmlInit, + nvmlShutdown, +) + +logger = init_logger(__name__) + + +def is_process_scoped_memory_available() -> bool: + """Check if NVML process-scoped memory tracking is available. + + When True, concurrent stage initialization is safe because each + process can accurately measure its own GPU memory via NVML. + When False, sequential initialization (file locks) is needed. + """ + try: + nvmlInit() + nvmlShutdown() + return True + except Exception: + return False + + +def parse_cuda_visible_devices() -> list[str | int]: + """Parse CUDA_VISIBLE_DEVICES into a list of device identifiers. + + Returns list of integers (physical indices) or strings (UUIDs/MIG IDs). + """ + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible_devices: + return [] + + result: list[str | int] = [] + for item in visible_devices.split(","): + item = item.strip() + if not item: + continue + try: + result.append(int(item)) + except ValueError: + # UUID (GPU-xxx) or MIG ID (MIG-xxx) + result.append(item) + return result + + +def get_device_handle(device_id: str | int): + """Get NVML device handle by index or UUID.""" + if isinstance(device_id, int): + return nvmlDeviceGetHandleByIndex(device_id) + else: + from vllm.third_party.pynvml import nvmlDeviceGetHandleByUUID + + return nvmlDeviceGetHandleByUUID(device_id) + + +def get_process_gpu_memory(local_rank: int) -> int | None: + """Get GPU memory used by current process via pynvml. + + Supports CUDA_VISIBLE_DEVICES with integer indices, UUIDs, or MIG IDs. + + Returns: + Memory in bytes used by this process, or None if NVML unavailable. + + Raises: + RuntimeError: If device validation fails (invalid index or UUID). + """ + from vllm.third_party.pynvml import nvmlDeviceGetCount + + my_pid = os.getpid() + visible_devices = parse_cuda_visible_devices() + + try: + nvmlInit() + except Exception as e: + logger.warning("NVML init failed, will use profiling fallback: %s", e) + return None + + try: + if visible_devices and local_rank < len(visible_devices): + device_id = visible_devices[local_rank] + try: + handle = get_device_handle(device_id) + except Exception as e: + raise RuntimeError( + f"Failed to get NVML handle for device '{device_id}' (local_rank={local_rank}). " + f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + ) from e + else: + # No CUDA_VISIBLE_DEVICES or local_rank out of range: use index directly + device_count = nvmlDeviceGetCount() + if local_rank >= device_count: + raise RuntimeError( + f"Invalid GPU device {local_rank}. Only {device_count} GPU(s) available. " + f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + ) + handle = nvmlDeviceGetHandleByIndex(local_rank) + + for proc in nvmlDeviceGetComputeRunningProcesses(handle): + if proc.pid == my_pid: + return proc.usedGpuMemory + return 0 + except RuntimeError: + raise + except Exception as e: + logger.warning("NVML query failed, will use profiling fallback: %s", e) + return None + finally: + try: + nvmlShutdown() + except Exception: + pass