[Hardware][Neuron] Refactor neuron support (#3471)

2024-03-21 18:22:17 -07:00
parent ea5f14e6ff
commit e90fc21f2e
33 changed files with 615 additions and 549 deletions
--- a/tests/spec_decode/utils.py
+++ b/tests/spec_decode/utils.py
@@ -123,7 +123,7 @@ def create_worker(cls: type,
        is_driver_worker=is_driver_worker,
    )

-    worker.init_model()
+    worker.init_device()
    worker.load_model()

    cache_config.num_gpu_blocks = num_gpu_blocks