[CI/CD] add neuron docker and ci test scripts (#3571)

2024-04-18 15:26:01 -07:00
parent 87fa80c91f
commit cd2f63fb36
6 changed files with 103 additions and 4 deletions
--- a/vllm/engine/async_llm_engine.py
+++ b/vllm/engine/async_llm_engine.py
@@ -335,8 +335,8 @@ class AsyncLLMEngine:
        engine_config = engine_args.create_engine_config()

        if engine_config.device_config.device_type == "neuron":
-            raise NotImplementedError("Neuron is not supported for "
-                                      "async engine yet.")
+            from vllm.executor.neuron_executor import NeuronExecutorAsync
+            executor_class = NeuronExecutorAsync
        elif engine_config.parallel_config.worker_use_ray:
            initialize_ray_cluster(engine_config.parallel_config)
            from vllm.executor.ray_gpu_executor import RayGPUExecutorAsync