详细说明

vllm_npu_0.4.2版本(即:vLLM 0.4.2版本昇腾框架适配代码)中修改了attention、engine、executor、model_executor、usage、worker六个模块,与vLLM原生框架中的同名模块一一对应进行热替换适配。vllm_npu_0.4.2版本可参考参考代码制作。

图1 vLLM 0.4.2版本架构图
表1 各模块修改内容介绍

模块

简介

attention模块

该模块重写了vLLM框架中的AttentionBackend类,对昇腾环境下对接MindIE LLM所需要的attention计算数据以及KV Cache的shape等关键信息进行了定义,并在该模块的初始化文件中对原框架的get_atten_backend函数进行了热替换,从而使得框架在昇腾环境下会运行该模块中的attention后端类。

engine模块

该模块重写了vLLM引擎的from_engine_args函数,vLLM 0.4.2版本中引擎会通过该类方法进行实例化,并在其中根据运行环境信息选择对应的不同executor;这里我们加入了新的判断逻辑分支,当检测到运行在昇腾环境下时,引擎会选择vllm_npu补丁包中定义的AscendExecutor和RayAscendExecutor类分别去进行单卡推理和多卡推理。另外,这里对vLLM原生框架的离线同步推理引擎LLMEngine和在线异步推理引擎AsyncLLMEngine的from_engine_args函数分别进行了重写替换,替换操作发生在该模块的初始化文件中。

executor模块

该模块中主要实现了四个executor类,其中AscendExecutor和AscendExecutorAsync用于单卡环境的同步和异步调用模式下的推理,RayAscendExecutor和RayAscendExecutorAsync用于多卡ray分布式环境的同步和异步调用模式下的推理。

此外,在ray_utils.py中对initialize_ray_cluster函数进行了重写,主要是因为在昇腾的npu环境下ray无法自动识别到npu的数量,因此需要手动指定。

model_executor模块

该模块为实际对接MindIE LLM模型推理与后处理的位置,其中包括layers模块和models模块,分别对应后处理和模型推理。

  • layers模块:该模块编写实现了AscendSampler类,进行vLLM原生框架的数据结构与模型仓底层数据结构之间的对接工作;具体实现代码如•layers模块实现代码:所示。
  • models模块:该模块编写实现了MindIELlmWrapper类,在该类中会对MindIE LLM提供的GeneratorTorch统一接口进行实例化操作,并从vLLM原生框架的数据结构中拆解出MindIE LLM所需要的模型推理参数,从而传给统一接口调用模型推理服务;另外,在进行warmup操作时使用的fake data构造操作也在该类中实现;具体实现代码如•layers模块实现代码:所示。

usage模块

该模块中的UsageMessage类的_report_usage_once成员函数进行了重写,修改了其中的torch.cuda.get_device_properties函数的使用方式,该函数目前在昇腾环境上的使用方式和GPU环境上有所差异。

worker模块

该模块实现了AscendWorker类,以供executor模块中的executor类进行调用;实现了AscendModelRunner类在AscendWorker中进行调用。

替换原生框架中CacheEngine的_allocate_kv_cache函数,主要是对生成kv_cache的数据格式进行了修改,从Torch.tensor修改为Tuple[torch.Tensor, torch.Tensor]。

AscendModelRunner类继承自原生框架中ModelRunner类,主要是为了对原生的load_model,execute_model和profile_run函数进行重写:vLLM新版本中执行模型调用时分为了先调用模型生成hidden_states,再使用一个process处理hidden_states得到logits,再进行最后的sample操作得到结果;而在MindIE模型仓中前两步操作是通过模型调用一步完成的,因此在这里进行了修改;profile_run函数的修改主要是为了构造warmup时使用的fake data。

除了上述的六个模块的适配外,还有一些主模块外的Python文件里的函数需要热替换,包括config.py中的DeviceConfig类中我们引入了NPU作为device_type,在utils.py文件中我们引入了is_ascend()函数用于检测当前运行环境是否为昇腾环境。最后,在npu_adaptor.py,我们对vLLM原框架中导入的一些昇腾环境下不具备的包(例如预编译的cuda算子、triton等)进行了屏蔽操作。