Merge origin/main into dev/vllm-align; resolve omni_stage.py conflicts

Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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?
 

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 "<optional quality filter>" \
   --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 "<optional quality filter>" \
   --height 480 \
   --width 832 \

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.

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 "<optional quality filter>" \
   --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 "<optional quality filter>" \
   --height 480 \
   --width 832 \

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"

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."
+    )

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

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]
 
 

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
 

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

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,156 +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
-    stage_engine = None
-    engine_init_error = None
-    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)
@@ -1262,25 +1117,6 @@ async def _stage_worker_async(
                 usage_context=usage_context,
                 engine_args=omni_engine_args,
             )
-    except Exception as e:
-        engine_init_error = e
-        logger.exception("[Stage-%s] Engine initialization failed: %s", stage_id, e)
-    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
-    # If engine init failed, re-raise the exception
-    if engine_init_error is not None:
-        raise engine_init_error
-    if stage_engine is None:
-        raise RuntimeError(f"[Stage-{stage_id}] Engine was not initialized")
     omni_stage.set_async_engine(stage_engine)
     if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats:
 

vllm_omni/entrypoints/openai/api_server.py

@@ -970,7 +970,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}"
 
@@ -1128,7 +1130,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())}"

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()
 
 

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
 

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)

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

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:

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