基于 MUSA 的大语言模型推理和服务框架vLLM

1. 引言​

vLLM是一个高性能且内存高效的大语言模型推理和服务框架，也是当前业界使用范围最广的大模型推理框架，截至目前github star数28.4k。该框架性能优秀，而且部署容易，使用CUDA/ROCm提供GPU加速能力。但vLLM目前不支持使用摩尔线程GPU进行加速，应广大摩尔线程客户及MUSA开发者的呼声，我们对该框架进行了适配。

2. vLLM与MUSA​

摩尔线程致力于构建完善好用的国产GPU应用生态，自主研发了MUSA架构及软件平台。现有的vLLM代码不支持摩尔线程GPU作为后端，因此我们新增了MUSA设备后端，从而让vLLM在摩尔线程GPU上流畅运行。

另外MUSA的一大优势是CUDA兼容，通过musify工具，我们可以快速将官方代码移植至MUSA软件栈，用户可以根据文档自行升级vLLM版本并适配MUSA软件栈。接下来我们将一步步介绍如何快速将vLLM适配到MUSA软件栈。

3. 软硬件依赖​

以下是MUSA开发人员适配时所使用的环境，作为参考：

• ubuntu20.04
• vLLM: v0.4.2
• MUSA SDK: rc3.0.1
• pytorch: v2.2.0
• torch_musa: v1.3.0
• GPU: 摩尔线程S4000
• 模型: meta-llama/Llama-2-7b-hf

4. MUSA适配​

4.1 MUSA移植​

通过使用musify工具，用户可以快速将原有的CUDA代码无缝迁移到MUSA软件栈，大大提升了用户在MUSA软件栈上开发的效率。musify工具是一个文本替换工具，用于将用户代码中CUDA相关的接口转换为MUSA的对应接口，然后使用MUSA软件栈下的mcc编译器编译成为摩尔线程GPU的可执行文件。

目前用户有两种途径使用musify工具： 一种是通过MUSA SDK中自带的工具： /usr/local/musa/tools/musify-text; 另一种是通过torch_musa中的torch_musa.utils.simple_porting模块; 这两种方式都是musify工具的入口，用户可按需使用。
musify工具提供了常见的接口转换的映射关系，用户一般只需要运行：

# pip install ahocorapy
/usr/local/musa/tools/musify-text <source files/dir to be transformed>

由于用户代码的多样性，有时可能需要用户补充一些映射关系，这里展示了用户自定义映射关系转换的方式：

from torch_musa.utils.simple_porting import SimplePorting

SimplePorting(cuda_dir_path="./csrc", mapping_rule={
    "#include <ATen/cuda/CUDAContext.h>": "#include \"torch_musa/csrc/aten/musa/MUSAContext.h\"",
    "#include <c10/cuda/CUDAGuard.h>": "#include \"torch_musa/csrc/core/MUSAGuard.h\"",
    "#include <ATen/cuda/Exceptions.h>": "#include \"torch_musa/csrc/core/MUSAException.h\"",
    "#include <c10/cuda/CUDAStream.h>": "#include \"torch_musa/csrc/core/MUSAStream.h\"",
    "at::kCUDA": "at::musa::kMUSA",
    "at::cuda::getCurrentCUDAStream()": "at::musa::getCurrentMUSAStream()",
    "__nv_bfloat16": "__mt_bfloat16",
    }).run()

限于篇幅，这里只展示了部分vLLM框架进行MUSA移植的映射关系，详情见： musa_porting.py

另外需要注意的是：CMakeLists.txt文件中有时会将源文件一个个添加进去，所以也需要修改其中的文件名及后缀(或者使用musify工具)。

4.2 添加MUSA后端​

该部分需要仿照代码中的CUDA后端，新增一个MUSA后端。该部分需要改动python层的代码，这里选取几个典型改动作为示例。

4.2.1 setup.py​

首先需要导入torch_musa库，同时使用torch_musa中的MUSAExtension将源文件添加到mcc的编译列表里，如下：

import torch_musa
from torch_musa.utils.musa_extension import BuildExtension, MUSAExtension

ext_modules = []
ext_modules.append(
        MUSAExtension(
            name="vllm_C",
            sources=[
                "csrc_musa/cache_kernels.mu",
                "csrc_musa/attention/attention_kernels.mu",
                "csrc_musa/pos_encoding_kernels.mu",
                "csrc_musa/activation_kernels.mu",
                "csrc_musa/layernorm_kernels.mu",
                "csrc_musa/musa_utils_kernels.mu",
                "csrc_musa/moe_align_block_size_kernels.mu",
                "csrc_musa/pybind.cpp",
                "csrc_musa/custom_all_reduce.mu",
            ],
            extra_compile_args= {"cxx": ['-O3'，'-std=c++17'],}
        )
    )

同时修改或新增一些后端的判断逻辑，让vLLM可以识别MUSA后端。
如在vllm/engine/arg_utils.py中修改为：

parser.add_argument("--device",
                    type=str,
                    default=EngineArgs.device,
                    choices=["auto", "cuda", "neuron", "cpu", "musa"],
                    help='Device type for vLLM execution.')

4.2.2 torch.musa​

vLLM框架中默认使用CUDA作为后端，因此在代码中直接调用了大量torch.cuda相关的接口，对应的需要修改成torch.musa相应的接口。或者添加判断条件，当后端设置为MUSA时，调用torch.musa相应的接口。如：

# 官方代码：
# device_name = torch.cuda.get_device_name().replace(" ", "_")
# 修改为MUSA接口：
device_name = torch.musa.get_device_name().replace(" ", "_")

4.2.3 FlashAttention​

摩尔线程 S4000(对应计算能力为mp_22)及之后的显卡均支持FlashAttention，对于pytorch框架我们需要使用scaled_dot_product_attention接口。因此我们需要在vllm/attention/backends/flash_attn.py文件中做如下修改：

import torch_musa
from torch.nn.functional import scaled_dot_product_attention

# enable musa flash attention
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_math_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(True)

attn_output = scaled_dot_product_attention(
                    query.contiguous(),
                    key.contiguous(),
                    value.contiguous(),
                    attn_mask=att_mask.contiguous(),
                    dropout_p=0.0,
                    is_causal=False,)

这样我们就可以体验到MUSA软件栈的深度学习加速库：muDNN，带来的FlashAttention加速，充分释放硬件的计算能力。

4.2.4 分布式​

CUDA架构使用NCCL作为分布式加速库，对应地，MUSA架构使用MCCL作为分布式加速库。我们需要在vllm/distributed/device_communicators/目录下，仿照官方的pynccl.py和pynccl_utils.py创建相应的通信库组件，然后将其中cuda，nccl字样分别替换为musa，mccl即可，然后在vllm/distributed/parallel_state.py中调用pymccl_utils模块，在摩尔线程GPU上使用MCCL进行分布式加速。

5. 示例​

接下来，我们通过一个简单的示例，展示下如何使用vLLM-MUSA进行大语言模型推理：

import torch
import torch_musa
from vllm import LLM, SamplingParams

# modify to your model path
model_path = "/workspace/models/Llama-2-7b-chat-hf/"

# prompts example
prompts = [
    "Hello, my name is",
    "The capital of France is",
]

# init vllm
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(model=model_path, trust_remote_code=True, device="musa")

# generate result
outputs = llm.generate(prompts, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

可以看到相对于原生仓库，我们仅需要导入torch_musa并将后端设置为musa。
至此，我们成功完成了vLLM在MUSA平台的适配工作。

6. 寄语​

vLLM-MUSA已经可在github获取: vllm_musa。

vLLM在MUSA平台快速便捷的适配过程，彰显了MUSA对CUDA的优良兼容性，助力用户业务实现快速高效迁移。我们期待更多的开发者，与我们一起，共同丰富完善MUSA开源社区，让MUSA变得越来越好。