vLLM 0.3.3 版本昇腾框架适配说明

vllm_npu_0.3.3版本的目录结构如下所示：

.
├── cover
│   ├── vllm
│   │   └── __init__.py
│   ├── requirements.txt
│   └── setup.py
├── examples
│   ├── start_server.sh
│   ├── test_offline.py
│   └── test_offline.sh
├── install.sh
└── vllm_npu
    ├── requirements.txt
    ├── setup.py
    ├── tests
    │   ├── models
    │   │   ├── __init__.py
    │   │   └── test_models.py
    │   └── sampler
    │       └── test_sampler.py
    └── vllm_npu
        ├── config.py
        ├── core
        │   ├── __init__.py
        │   └── scheduler.py
        ├── engine
        │   ├── __init__.py
        │   ├── llm_engine.py
        │   └── ray_utils.py
        ├── __init__.py
        ├── model_executor
        │   ├── ascend_model_loader.py
        │   ├── __init__.py
        │   ├── layers
        │   │   ├── __init__.py
        │   │   └── sampler.py
        │   ├── models
        │   │   ├── ascend
        │   │   │   ├── __init__.py
        │   │   │   └── mindie_llm_wrapper.py
        │   │   └── __init__.py
        │   └── utils.py
        ├── npu_adaptor.py
        ├── utils.py
        └── worker
            ├── ascend_worker.py
            ├── cache_engine.py
            ├── __init__.py
            └── model_runner.py

vllm_npu_0.3.3版本中主要修改了core、engine、model_executor、worker四个模块，与vllm原生框架中的同名模块一一对应进行热替换适配。

表1 各模块修改内容介绍
模块	简介
core	该模块中对于 “num_batched_tokens”的计算与后端有所不同，MindIE后端的实现方法与0.4.2版本一致，因此在此处需要修改该变量计算逻辑。
engine	该模块替换了部分内容，主要是为了确保vLLM框架可以正确获取到后续的worker与model_executor模块。例如：在ray模块启动时新增指定 “num_gpus”参数为 “parallel_config.world_size”，解决ray模块启动问题；并且在DEVICE_TO_WORKER_MODULE_MAP常量字典中新增键值对："npu"、 "vllm_npu.worker.ascend_worker"，确保在启动worker时可以正确使用ascend_worker。
model_executor	该模块负责与MindIE LLM模型的推理与后处理进行对接，包含两个子模块：“layers”和 “models”。其中，“layers” 模块主要负责后处理操作，而“models”模块则专注于模型推理。在“models”模块中，编写了一个名为 “MindIELlmWrapper”的类。这个类用于实例化 MindIE LLM提供的 “GeneratorTorch”统一接口，并从vLLM 原生框架的数据结构中提取出MindIE LLM所需的模型推理参数，进而调用统一接口进行模型推理。此外，该类还实现了在进行预热（warmup）操作时所需的虚拟数据的构造过程。同时添加了ascend_model_loader.py文件，为后续的“ModelRunner”加载模型提供“get_model”方法。
worker	该模块覆写了“Worker”类与“ModelRunner”类中的部分函数，在ascend_worker.py和model_runner.py中与 “MindIELlmWrapper”交互，确保模型的加载、预热、推理和后处理等过程能够正确在后端执行。此外，修改了原生框架中“CacheEngine”的“_allocate_kv_cache”函数，确保与后端接收的参数形式一致。

推理流程与部分关键代码的说明

在vLLM中从启动引擎到完成推理，有如下步骤：

输入模型架构名和权重路径：首先提供模型的基本信息。
初始化引擎参数：
“model str”-> “EngineArgs()”->“LLMEngine.from_engine_args()”->“LLMEngine.init()”
配置模型：
“ModelConfig()”：如果设置了“VLLM_USE_MODELSCOPE”变量，则从 ModelScope hub 下载对应的模型，并设置下载路径为模型权重路径。

初始化工作器：“LLMEngine._init_workers()”->“Worker.init_model()”：在这一步骤中，设置NPU环境、分布式配置和随机种子等。“Worker.load_model()”->“ModelRunner.load_model()”->“vllm.model_executor.model_loader.get_model()”

    def load_model(self):
        # Add an attribute dynamically, not graceful, but useful
        self.model_config.mindie_config = {
            "backend_type": BackendType.ATB,
            "model_id": self.model_config.model,
            "rank": self.rank,
            "local_rank": self.local_rank,
            "world_size": self.parallel_config.world_size,
            "npu_device_id": self.local_rank,
        }
        self.model_runner.load_model()
        del self.model_config.mindie_config

加载模型：

“ascend_model_loader.get_model()”->“mindie_llm_wrapper.load_weights()” ->“mindie_llm.text_generator.adapter.GeneratorTorch(model_config)”，由MindIE LLM实现模型加载的逻辑。

    def load_weights(self, load_format: Literal["auto", "safetensors", "pt"] = "auto"):
        """Load model weights.
        Args:
            load_format (Literal[auto, safetensors, pt], optional): The format of weights. Defaults to "auto".
        Raises:
            ValueError: Format not supported.
        """
        if load_format not in ["auto", "safetensors", "pt"]:
            raise ValueError("load-format support [safetensors, pt]")
        self.weight_dtype = torch.get_default_dtype()
        torch.set_default_dtype(torch.float32)
        # Replacing the deprecated method with GeneratorTorch
        self.atb_model = GeneratorTorch(self.mindie_config)
        self.sampler = Sampler(self.atb_model)
        self.sampler.logits_as_hidden_states = True
        torch.set_default_dtype(self.weight_dtype)

创建模型实例：
回到“model_loader.get_model()”函数中根据类创建模型实例。如果存在量化方法，也在此实例化。

实际推理：

“LLMEngine.step”->“worker.execute_model”-> “model_runner.execute_model”->“mindie_llm_wrapper.foward” -> MindIE后端。其中在 “model_runner”层还需要在“prepare_input_tensor”及其私有方法中对输入进来的数据进行处理。

同样，后处理部分也是由MindIE后端来进行处理与看护。

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[KVCache],
        input_metadata: InputMetadata,
        lora_requests: List[LoRARequest],
    ) -> torch.Tensor:
        if kv_caches[0][0] is None:
            kv_caches, block_tables, slots = self.create_dummy_kv_cache(input_metadata, input_ids)
        else:
            if input_metadata.is_prompt:
                block_tables = torch.tensor([0], dtype=torch.int32, device="npu")
            else:
                block_tables = input_metadata.block_tables
            slots = input_metadata.slot_mapping
        if input_metadata.is_prompt:
            input_lengths = input_metadata.prompt_lens.to(torch.int32)
            max_seq_len = int(input_metadata.prompt_lens.max())
            lm_head_indices = (input_metadata.prompt_lens.cumsum(dim=-1) - 1).to(torch.int64)
        else:
            input_lengths = input_metadata.context_lens
            max_seq_len = input_metadata.max_context_len
            lm_head_indices = None
        adapter_ids = [lora_request.lora_name if lora_request else None for lora_request in lora_requests]
        logits = self.atb_model.forward_tensor(
            input_ids,
            positions,
            input_metadata.is_prompt,
            kv_caches,
            block_tables,
            slots,
            input_lengths,
            max_seq_len,
            lm_head_indices,
            adapter_ids=adapter_ids,  # batch_size len
        )
        return logits
    def sample(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        # hidden_states is logits
        next_tokens = self.sampler(hidden_states, sampling_metadata)
        return next_tokens

对应修改量

表2 对应修改量
重要组件	代码量	作用
Scheduler	~200行	修改“num_batched_tokens”计算逻辑。
Cache Engine	~30行	调整KV Cache参数结构，兼容底层接口。
Worker	~250行	构建后端统一接口所需的配置参数，并且初始化昇腾硬件环境。
Model Loader	~50行	重新编写get_model函数，对接下层调用。
Model Runner	~300行	传递配置参数，替换get_model，修改profile_run中对KV Cache的初始化方式，修改LoRA对应参数传递方式。
MindIE LLM Wrapper	~150行	负责加载、生成虚拟KV Cache、推理、后处理操作与后端的对接。
Sampler	~300行	从SamplingMetadata重构参数，传递至后端Sampler，替代vLLM自有后处理方式，并将结果恢复至vLLM结构。

除此之外，还有少部分文件需要修改，例如：config.py和utils.py等文件需要修改以解决少量在运行时出现的问题。

父主题： 适配说明