neuron.py 1.85 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from typing import TYPE_CHECKING, Optional
4

5
6
from vllm.logger import init_logger

7
8
from .interface import Platform, PlatformEnum

9
10
11
12
13
if TYPE_CHECKING:
    from vllm.config import VllmConfig
else:
    VllmConfig = None

14
15
logger = init_logger(__name__)

16
17
18

class NeuronPlatform(Platform):
    _enum = PlatformEnum.NEURON
19
    device_name: str = "neuron"
20
    device_type: str = "neuron"
21
    ray_device_key: str = "neuron_cores"
22
    supported_quantization: list[str] = ["neuron_quant"]
23
    device_control_env_var: str = "NEURON_RT_VISIBLE_CORES"
24
25
26
27

    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        return "neuron"
28

29
30
31
32
    @classmethod
    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
        return False

33
34
35
36
37
38
    @classmethod
    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
        parallel_config = vllm_config.parallel_config
        if parallel_config.worker_cls == "auto":
            parallel_config.worker_cls = \
                "vllm.worker.neuron_worker.NeuronWorker"
39

40
41
42
        if parallel_config.world_size > 1:
            parallel_config.distributed_executor_backend = "uni"

43
44
        assert (vllm_config.lora_config
                is None), "LoRA is not supported for Neuron backend."
45
46
47
        assert (not vllm_config.speculative_config
                ), "Speculative decoding not yet supported for Neuron backend."

48
49
50
51
52
53
        cache_config = vllm_config.cache_config
        if cache_config:
            # neuron needs block_size = max_model_len
            vllm_config.cache_config.block_size = \
                vllm_config.model_config.max_model_len

54
55
56
57
    @classmethod
    def is_pin_memory_available(cls) -> bool:
        logger.warning("Pin memory is not supported on Neuron.")
        return False
58
59
60
61

    @classmethod
    def use_all_gather(cls) -> bool:
        return True