From 3a508d7b7a250c7c6a05b232bb4c5a7f71114c2c Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 16 Mar 2026 02:32:39 +0000 Subject: [PATCH] rebase: align vllm-omni with vLLM e9163b536e72 Signed-off-by: tzhouam --- docker/Dockerfile.ci | 2 +- tests/entrypoints/test_omni_diffusion.py | 20 +++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 1 + .../core/sched/omni_generation_scheduler.py | 1 + vllm_omni/entrypoints/cli/benchmark/main.py | 2 +- vllm_omni/entrypoints/openai/api_server.py | 35 ++++++++++++++----- .../models/mimo_audio/mimo_audio_llm.py | 2 -- .../qwen2_5_omni/qwen2_5_omni_talker.py | 2 -- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 13 ++++--- .../qwen3_omni/qwen3_omni_moe_thinker.py | 20 +++++++---- vllm_omni/worker/gpu_ar_model_runner.py | 23 ++++++------ vllm_omni/worker/gpu_ar_worker.py | 6 ++-- .../worker/gpu_generation_model_runner.py | 10 +++--- vllm_omni/worker/gpu_generation_worker.py | 6 ++-- 14 files changed, 97 insertions(+), 46 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 117aba153..6e7b64edd 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -17,7 +17,7 @@ RUN apt-get update && \ # Note: --prerelease=allow is required because nightly/commit wheels use # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default, # causing it to fall back to the older stable release from PyPI. -RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/84e436ed1c94b1b94f809927b5d6bff45f7af919 +RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/e9163b536e721c431500f6f43ace22fcb3532e7e RUN uv pip install --system --no-cache-dir ".[dev]" diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index 9e555fa85..ea4402ebd 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -1371,9 +1371,19 @@ def test_dtype_normalization_valid_types( _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) + monkeypatch.setattr( + "vllm_omni.entrypoints.omni.omni_snapshot_download", + lambda model_id: model_id, + raising=False, + ) + + import vllm_omni.entrypoints.omni as omni_module from vllm_omni.entrypoints.omni import Omni + monkeypatch.setattr(omni_module, "omni_snapshot_download", lambda model_id: model_id) + monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", _fake_loader, @@ -1417,9 +1427,19 @@ def test_dtype_normalization_invalid_types(monkeypatch, mocker: MockerFixture, f _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) + _setup_connector_mocks(monkeypatch, mocker) + monkeypatch.setattr( + "vllm_omni.entrypoints.omni.omni_snapshot_download", + lambda model_id: model_id, + raising=False, + ) + + import vllm_omni.entrypoints.omni as omni_module from vllm_omni.entrypoints.omni import Omni + monkeypatch.setattr(omni_module, "omni_snapshot_download", lambda model_id: model_id) + monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", _fake_loader, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 540fc391c..3629ade63 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -364,6 +364,7 @@ class OmniARScheduler(VLLMScheduler): if stopped_preempted_reqs: # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) + self.skipped_waiting.remove_requests(stopped_preempted_reqs) # [Main] Handle failed KV load requests if failed_kv_load_req_ids and not self.recompute_kv_load_failures: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index b2e64224b..64522b909 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -502,6 +502,7 @@ class OmniGenerationScheduler(VLLMScheduler): if stopped_preempted_reqs: # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) + self.skipped_waiting.remove_requests(stopped_preempted_reqs) # Handle failed KV load requests if failed_kv_load_req_ids and not self.recompute_kv_load_failures: diff --git a/vllm_omni/entrypoints/cli/benchmark/main.py b/vllm_omni/entrypoints/cli/benchmark/main.py index 8880e35c7..865064d1e 100644 --- a/vllm_omni/entrypoints/cli/benchmark/main.py +++ b/vllm_omni/entrypoints/cli/benchmark/main.py @@ -9,7 +9,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase if typing.TYPE_CHECKING: - from vllm.utils import FlexibleArgumentParser + from vllm.utils.argparse_utils import FlexibleArgumentParser class OmniBenchmarkSubcommand(CLISubcommand): diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 4929f951c..c83cb7b7b 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -35,13 +35,6 @@ from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server - -# vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. -# Keep a fallback for older/newer upstream layouts during rebase windows. -try: - from vllm.entrypoints.serve.instrumentator.basic import base -except ModuleNotFoundError: - from vllm.entrypoints.openai.basic.api_router import base from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -68,10 +61,15 @@ from vllm.entrypoints.openai.speech_to_text.serving import ( ) from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.classify.serving import ServingClassification -from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding +from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling from vllm.entrypoints.pooling.score.serving import ServingScores from vllm.entrypoints.serve.disagg.serving import ServingTokens + +# vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. +# Keep a fallback for older/newer upstream layouts during rebase windows. +from vllm.entrypoints.serve.instrumentator.basic import base +from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization from vllm.entrypoints.utils import ( load_aware_call, @@ -306,6 +304,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ssl_certfile=args.ssl_certfile, ssl_ca_certs=args.ssl_ca_certs, ssl_cert_reqs=args.ssl_cert_reqs, + ssl_ciphers=args.ssl_ciphers, h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, h11_max_header_count=args.h11_max_header_count, **uvicorn_kwargs, @@ -579,6 +578,22 @@ async def omni_init_app_state( ) await state.openai_serving_models.init_static_loras() + state.openai_serving_render = OpenAIServingRender( + model_config=engine_client.model_config, + renderer=engine_client.renderer, + io_processor=engine_client.io_processor, + model_registry=state.openai_serving_models.registry, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, + ) + state.openai_serving_responses = ( OpenAIServingResponses( engine_client, @@ -603,6 +618,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, + openai_serving_render=state.openai_serving_render, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -629,6 +645,7 @@ async def omni_init_app_state( OpenAIServingCompletion( engine_client, state.openai_serving_models, + openai_serving_render=state.openai_serving_render, request_logger=request_logger, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, @@ -691,6 +708,7 @@ async def omni_init_app_state( request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, trust_request_chat_template=args.trust_request_chat_template, ) state.openai_serving_transcription = ( @@ -718,6 +736,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, + openai_serving_render=state.openai_serving_render, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index c1b71e076..1424ca775 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -773,7 +773,6 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, - handle_oov_mm_token: bool = False, ) -> torch.Tensor: # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: @@ -783,7 +782,6 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) def base_local_forward( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 210d262f1..dbb6c49ef 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -109,7 +109,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, - handle_oov_mm_token: bool = False, ) -> torch.Tensor: # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: @@ -119,7 +118,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) def forward( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 4560ede39..5c69fca66 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -116,6 +116,15 @@ class Qwen2_5OmniThinkerMultiModalProcessor( tokenizer = self.info.get_tokenizer() audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>") use_audio_in_video = audio_pad_id not in prompt_ids + # for mutilmodality cache + if any(item is None for item in mm_kwargs["video"]): + video_token_id = self.info.get_hf_config().video_token_id + audio_token_id = self.info.get_hf_config().audio_token_id + video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) + audio_updates_num = len(mm_prompt_updates.get("audio", [])) + video_updates_num = len(mm_prompt_updates.get("video", [])) + if video_audio_item_num != video_updates_num + audio_updates_num: + use_audio_in_video = True if is_update_applied: mm_placeholders = self._find_mm_placeholders( @@ -600,7 +609,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, - handle_oov_mm_token: bool = False, ) -> torch.Tensor: if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) @@ -609,7 +617,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) if len(multimodal_embeddings) == 0: @@ -632,7 +639,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, self.get_language_model().embed_input_ids, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) return merge_interleaved_embeddings( inputs_embeds, @@ -649,7 +655,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) def forward( diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index ca7bb58c7..5111fe1d4 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -146,9 +146,11 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32) cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens]) - sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(self.attn_backend, cu_seqlens) - if sequence_lengths is not None: - sequence_lengths = torch.from_numpy(sequence_lengths).to(self.device, non_blocking=True) + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( + self.attn_backend, + cu_seqlens, + self.device, + ) max_seqlen = torch.tensor( MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens), dtype=torch.int32, @@ -475,6 +477,15 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( tokenizer = self.info.get_tokenizer() audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>") use_audio_in_video = audio_pad_id not in prompt_ids + # for mutilmodality cache + if any(item is None for item in mm_kwargs["video"]): + video_token_id = self.info.get_hf_config().video_token_id + audio_token_id = self.info.get_hf_config().audio_token_id + video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) + audio_updates_num = len(mm_prompt_updates.get("audio", [])) + video_updates_num = len(mm_prompt_updates.get("video", [])) + if video_audio_item_num != video_updates_num + audio_updates_num: + use_audio_in_video = True # normal case with `use_audio_in_video=False` if is_update_applied: @@ -966,13 +977,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( multimodal_embeddings: MultiModalEmbeddings | None = None, *, is_multimodal: torch.Tensor | None = None, - handle_oov_mm_token: bool = False, ) -> torch.Tensor: inputs_embeds = self._embed_text_input_ids( input_ids, self.language_model.embed_input_ids, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: @@ -1064,7 +1073,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal, - handle_oov_mm_token=handle_oov_mm_token, ) def forward( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d7d45031a..166cc6231 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -123,7 +123,7 @@ class GPUARModelRunner(OmniGPUModelRunner): request_id_resolver=self._resolve_global_request_id, ) - if self.vllm_config.model_config.enable_return_routed_experts: + if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() if capturer is not None: capturer.clear_buffer() # noqa @@ -141,7 +141,7 @@ class GPUARModelRunner(OmniGPUModelRunner): # Update persistent batch states. self._update_states(scheduler_output) - if has_ec_transfer() and get_ec_transfer().is_producer: + if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -276,9 +276,9 @@ class GPUARModelRunner(OmniGPUModelRunner): # Run the model. # Use persistent buffers for CUDA graphs. - # When spec decode is enabled, delay clearing connector metadata - # until after draft model runs in sample_tokens. - clear_kv_metadata = self.speculative_config is None + # When spec decode is enabled, defer connector finalization + # (wait_for_save + clear metadata) until after draft model runs. + defer_kv_connector_finalize = self.speculative_config is not None with ( set_forward_context( attn_metadata, @@ -292,7 +292,8 @@ class GPUARModelRunner(OmniGPUModelRunner): ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output( - scheduler_output, clear_metadata=clear_kv_metadata + scheduler_output, + defer_finalize=defer_kv_connector_finalize, ) as kv_connector_output, ): model_output = self._model_forward( @@ -533,11 +534,11 @@ class GPUARModelRunner(OmniGPUModelRunner): # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) - # Clear KV connector metadata after draft model runs (if spec decode). - # This was deferred from target model forward to allow draft model - # to also save its KV cache. + # Finalize KV connector (wait_for_save + clear metadata) after + # draft model runs. Deferred from target model forward to allow + # draft model to also save its KV cache. if self.speculative_config is not None: - self.clear_kv_connector_metadata() + self.finalize_kv_connector() with record_function_or_nullcontext("gpu_model_runner: eplb"): self.eplb_step() @@ -591,7 +592,7 @@ class GPUARModelRunner(OmniGPUModelRunner): payload.update(mm_payload) pooler_output.append(payload) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): - if self.model_config.enable_return_routed_experts: + if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() if capturer is not None: capturer.save_captured_experts(indices=self.slot_mapping) # noqa diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 8734524a2..4abe21964 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -48,17 +48,17 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK self.local_rank += dp_local_rank * tp_pp_world_size - assert self.local_rank < torch.cuda.device_count(), ( + assert self.local_rank < torch.accelerator.device_count(), ( f"DP adjusted local rank {self.local_rank} is out of bounds. " ) - visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0 assert self.parallel_config.local_world_size <= visible_device_count, ( f"local_world_size ({self.parallel_config.local_world_size}) must " f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) self.device = torch.device(f"cuda:{self.local_rank}") - current_platform.set_device(self.device) + torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 58e4dfd72..4db683a8b 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -87,7 +87,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - if self.vllm_config.model_config.enable_return_routed_experts: + if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() if capturer is not None: capturer.clear_buffer() # noqa @@ -265,9 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Run the model. # Use persistent buffers for CUDA graphs. - # When spec decode is enabled, defer clearing KV connector metadata - # until after draft model runs in sample_tokens. - clear_kv_metadata = self.speculative_config is None + # When spec decode is enabled, defer connector finalization + # (wait_for_save + clear metadata) until after draft model runs. + defer_kv_connector_finalize = self.speculative_config is not None with ( set_forward_context( attn_metadata, @@ -282,7 +282,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output( scheduler_output, - clear_metadata=clear_kv_metadata, + defer_finalize=defer_kv_connector_finalize, ) as kv_connector_output, ): outputs = self._run_generation_model( diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 341797ebb..267ed61c0 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -48,17 +48,17 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK self.local_rank += dp_local_rank * tp_pp_world_size - assert self.local_rank < torch.cuda.device_count(), ( + assert self.local_rank < torch.accelerator.device_count(), ( f"DP adjusted local rank {self.local_rank} is out of bounds. " ) - visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0 assert self.parallel_config.local_world_size <= visible_device_count, ( f"local_world_size ({self.parallel_config.local_world_size}) must " f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) self.device = torch.device(f"cuda:{self.local_rank}") - current_platform.set_device(self.device) + torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype)