[MUSA][Feat] torch.accelerator support (#3101)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>

vllm_omni/worker/gpu_generation_worker.py

@@ -28,7 +28,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
 
     @instrument(span_name="Init device")
     def init_device(self):
-        if self.device_config.device_type == "cuda":
+        if self.device_config.device_type in ("cuda", "musa"):
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
             parallel_config = self.parallel_config
@@ -51,7 +51,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
                 assert self.local_rank < torch.accelerator.device_count(), (
                     f"DP adjusted local rank {self.local_rank} is out of bounds. "
                 )
-                visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
+                visible_device_count = torch.accelerator.device_count()
                 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 "