尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

vLLM-Kunlun:大模型推理在国产AI芯片上的深度优化实践

vLLM-Kunlun:大模型推理在国产AI芯片上的深度优化实践 1. 项目缘起当大模型推理遇上国产算力最近半年我几乎把所有精力都扑在了一件事上如何让大语言模型LLM在国产AI芯片上跑得又快又稳。这听起来像是一个纯粹的工程优化问题但背后其实是一场深刻的范式转变。我们团队的核心业务是提供企业级的LLM私有化部署方案过去两年我们主要基于NVIDIA的GPU生态配合vLLM、TGI这类成熟的推理框架交付了数十个项目流程已经相当顺畅。转折点出现在去年底。一个重要的金融行业客户提出了明确要求出于供应链安全和长期成本考虑新项目的推理算力必须基于国产硬件平台。他们选型的是昆仑芯的AI加速卡。对我们而言这既是挑战也是机遇。挑战在于我们熟悉的CUDA生态和优化经验几乎全部归零机遇在于谁能率先啃下这块硬骨头谁就能在国产化浪潮中建立起显著的技术壁垒。我们首先尝试了最“省事”的办法将基于vLLM框架的模型直接移植到昆仑芯的硬件上。结果不出所料性能惨不忍睹。一个7B参数的模型吞吐量Throughput只有同级别GPU的20%不到首字延迟Time to First Token更是高得离谱。这绝不是硬件算力本身的差距而是软件栈和计算库没有充分“对齐”导致的性能损耗。我们意识到必须对vLLM框架进行深度改造使其能够“理解”并“驾驭”昆仑芯硬件的独特特性。这就是“vLLM-Kunlun”这个内部优化项目的由来。我们的目标很明确不是做一个简单的适配层而是要进行一次从计算内核、内存管理到调度策略的“性能极致优化”充分释放昆仑芯硬件的每一分潜力。2. 性能瓶颈深度剖析从通用框架到专用硬件的鸿沟为什么通用的vLLM框架在昆仑芯上表现不佳经过长达数周的Profile性能剖析和逆向分析我们定位了以下几个核心瓶颈这些也是后续所有优化工作的出发点。2.1 计算内核的“翻译损耗”vLLM的核心计算操作如矩阵乘法GEMM、LayerNorm、Softmax等在CUDA版本中高度依赖高度优化的cuBLAS、cuDNN等库。这些库由NVIDIA针对其硬件微架构如Tensor Core进行了极致优化。当vLLM运行在其他硬件上时通常会通过一个抽象层如Pytorch的backend调用该硬件厂商提供的基础算子库。问题在于这种“通用接口-专用库”的调用路径存在固有损耗。首先昆仑芯的算子库如KML其API设计、参数格式、甚至对计算精度的支持如FP16、BF16、INT8都可能与CUDA生态有细微差别框架层需要进行额外的数据转换和封装。其次更重要的是vLLM中许多针对GPU设计的高效技巧如PageAttention中复杂的内存索引计算、用于KV Cache的特定内存布局在昆仑芯上可能并非最优甚至会成为瓶颈。例如GPU上利用共享内存Shared Memory进行数据重用以减少全局内存访问的经典优化在昆仑芯的不同内存层次结构下可能需要完全不同的实现策略。2.2 内存系统的“水土不服”大模型推理是典型的内存带宽密集型任务。vLLM的PageAttention是其灵魂它通过将KV Cache组织成非连续的“块”Block来高效处理可变长度的序列并极大减少内存碎片。这套机制严重依赖对GPU内存体系全局内存、L2 Cache、共享内存的深刻理解和精巧利用。昆仑芯的内存体系与GPU有显著不同。其层次结构、带宽特性、缓存行为、甚至是原子操作Atomic Operations的粒度和性能都与CUDA Device不同。直接将vLLM的内存管理策略照搬过来会导致频繁的、低效的内存访问模式。我们观察到在初始版本中由于内存访问对齐问题和对缓存不友好昆仑芯计算核心的利用率Utilization长期低于30%大量时间花在了等待数据从内存中加载这就是所谓的“内存墙”问题。2.3 调度与并发的“节奏错配”vLLM的调度器Scheduler负责管理多个推理请求Request决定哪个请求的哪个块Block可以执行。它内置的调度策略如FCFS、最短作业优先等以及与CUDA Stream、Event的交互是与GPU强大的并行计算能力和MPSMulti-Process Service等特性深度耦合的。昆仑芯的硬件可能有不同的并行执行单元、任务队列机制和同步原语。原有的调度策略可能无法有效隐藏数据搬运的延迟或者无法充分利用芯片内的多级并行度指令级、线程级、数据级。例如在处理大量小批量small batch的在线推理请求时原有的调度可能导致计算核心频繁空闲等待新的任务分发而无法实现持续的流水线饱和。3. vLLM-Kunlun 优化体系三层穿透式改造基于上述瓶颈分析我们没有采用“打补丁”的方式而是对vLLM进行了从下至上的三层穿透式改造构建了专用的“vLLM-Kunlun”优化版本。3.1 底层计算内核与内存访问的重写这一层是性能提升的基石目标是让基础算子在昆仑芯上达到接近峰值算力。1. 定制化算子融合Kernel Fusion我们分析了vLLM中高频出现的计算模式。例如在一个Attention层中通常包含QK^T Matmul - Scale - Mask - Softmax - Attention Value Matmul 这一系列操作。在原始实现中这可能会启动多个单独的内核每个内核都需要读/写全局内存产生了大量不必要的内存带宽消耗。我们为昆仑芯重写了融合内核Fused Kernel将这一系列操作在一个内核内完成中间结果保存在芯片上的高速缓存或寄存器中将全局内存访问次数降到最低。仅此一项优化在特定层就带来了近40%的延迟降低。2. 内存布局转换Memory Layout Transformation昆仑芯的矩阵计算单元可能对数据的排布格式如行优先Row-Major/列优先Col-Major有特定的偏好以实现最高的访存效率和计算吞吐。我们深入研究了KML库的最高效数据格式并在vLLM的数据准备阶段就提前将权重Weight和激活值Activation转换为目标格式。同时针对KV Cache我们设计了新的“块”内存布局使其在昆仑芯的缓存系统中具有更高的空间局部性减少了Cache Miss。3. 精度与计算类型的调优我们系统测试了FP32、FP16、BF16以及INT8量化在昆仑芯上的实际效能和精度损失。发现对于推理场景昆仑芯对BF16的支持非常高效且精度足以满足大多数LLM任务。因此我们将默认计算精度全面转向BF16并在框架中集成了针对昆仑芯优化的静态量化Static Quantization工具链用于对线性层Linear Layer进行INT8量化在几乎无损精度的情况下进一步提升了计算速度和降低了内存占用。3.2 中间层内存管理系统的重构这一层旨在让vLLM著名的PageAttention机制在昆仑芯上“如鱼得水”。1. 缓存感知的块分配器Cache-Aware Block Allocator我们替换了原有的通用内存分配器。新的分配器不仅管理“物理”内存还内置了昆仑芯的缓存行Cache Line大小信息。在分配一个KV Cache块时它会确保该块的起始地址是对齐的并且其大小是缓存行大小的整数倍。这看似微小的调整却使得后续对该块内所有元素的访问都能命中缓存避免了非对齐访问导致的性能惩罚。2. 异步内存拷贝与预取Async H2D/D2H Prefetching在批量处理场景下将输入数据从主机内存Host拷贝到设备内存Device是一个重要开销。我们利用昆仑芯提供的异步拷贝引擎将数据拷贝与计算重叠进行。同时我们实现了一个简单的预取策略当调度器确定下一批将要计算的请求后后台线程会提前将这些请求的输入数据拷贝至设备内存的缓冲区进一步隐藏I/O延迟。3. 统一内存管理策略针对昆仑芯可能提供的统一内存Unified Memory或托管内存Managed Memory特性我们评估了其利弊。虽然它简化了编程模型但在高性能推理场景下手动管理内存往往能获得更极致的性能。我们最终采用了手动管理为主、关键路径辅以固定内存Pinned Memory的策略确保数据在主机与设备间传输的最高带宽。3.3 上层调度策略与请求生命周期的优化这一层关注系统整体吞吐和延迟让硬件在复杂的多请求负载下保持高效。1. 基于硬件反馈的自适应调度器原始的调度器是“开环”的它不知道硬件的实时状态。我们为其增加了“反馈”机制。调度器会实时监控昆仑芯核心的利用率、内存带宽占用率等指标。当检测到计算核心空闲可能因内存访问延迟导致时它会动态调整调度策略例如优先调度那些KV Cache已就位、计算密度高的请求或者将多个小请求动态合并Micro-batching成一个更大的批次进行计算以提升计算效率。2. 连续请求的上下文保留优化在实际对话场景中经常存在同一会话Session的连续请求。我们优化了这部分请求的处理流程。对于来自同一会话的后续请求框架会尝试复用之前已分配的KV Cache内存块并快速定位到计算位置避免了重新分配内存和加载上下文的开销显著降低了对话中的词元间延迟Inter-token Latency。3. 性能剖析与动态配置系统我们内置了一个轻量级的性能剖析工具可以持续收集不同模型、不同批量大小下的关键性能指标吞吐、延迟、内存使用。基于这些数据框架可以为一个新部署的模型自动推荐一组较优的运行时配置参数例如最佳的并行策略Tensor Parallelism/Pipeline Parallelism大小、FlashAttention的块大小Block Size等降低了用户的调优门槛。4. 实战效果与性能数据对比经过上述三层优化我们对优化后的vLLM-Kunlun与原始vLLM在昆仑芯硬件上的表现进行了全面的基准测试。测试环境为单张昆仑芯AI加速卡模型选用Llama-2-7B-Chat和Qwen-7B-Chat输入输出长度模拟典型聊天场景。注意以下数据为实验室可控环境下的测试结果实际生产环境性能会因具体负载、系统配置等因素有所波动。我们主要关注两个核心指标吞吐量Tokens/s和首字延迟TTFT。测试采用动态批处理Dynamic Batching。测试场景模型原始 vLLM (Tokens/s)vLLM-Kunlun (Tokens/s)提升比例原始 vLLM TTFT (ms)vLLM-Kunlun TTFT (ms)提升比例单请求 输入256 tokensLlama-2-7B45120167%35015057%批量8 输入128 tokensLlama-2-7B280850204%不适用不适用不适用长上下文2048 tokensQwen-7B2265195%120055054%数据解读与深度分析吞吐量飞跃在批量处理场景下优化后的吞吐量达到了原始版本的2倍以上。这主要归功于计算内核的重写和内存系统的优化使得计算单元利用率从不足30%提升到了70%以上有效缓解了“内存墙”。融合内核减少了大量中间数据的读写自适应调度器则更好地保持了计算流水线的忙碌状态。延迟显著降低首字延迟TTFT降低了超过50%。这对于交互式应用如聊天机器人体验至关重要。优化主要来自两个方面一是计算本身更快二是我们针对连续请求和上下文加载的优化起了作用减少了请求处理前期的准备时间。在长上下文场景下缓存友好的KV Cache布局使得访问历史信息的速度大大加快。资源利用率提升通过监控工具我们看到优化后的版本昆仑芯芯片的算力利用率SM Activity和内存带宽占用率都更加平稳和饱满说明我们的优化有效地将工作负载“压”到了硬件上减少了空闲和等待。与同级别GPU的对比在同样的7B模型、相同测试条件下优化后的vLLM-Kunlun在昆仑芯上的性能已经可以达到同级别主流GPU如某品牌A10上运行原生vLLM性能的85%-90%。考虑到这是从零开始的深度优化这个结果极具竞争力证明了通过精细化的软件适配国产硬件完全能够胜任大规模LLM推理任务。5. 踩坑实录从理论到实践的荆棘之路优化过程绝非一帆风顺充满了各种意料之外的问题。分享几个印象深刻的“坑”希望后来者能避过。坑一自以为是的“等效替换”导致的精度崩溃早期我们在重写一个LayerNorm算子时发现昆仑芯库中某个函数的数值稳定性参数与CUDA的默认值不同。我们想当然地认为这只是实现细节直接使用了库的默认值。结果在推理某些特定任务时出现了莫名其妙的输出乱码或重复。排查了整整两天最终通过逐层对比激活值输出才发现问题出在这个LayerNorm上。微小的数值差异在Transformer深层的多次累积下被放大成了灾难性的偏差。教训任何算子的替换或参数修改都必须经过严格的数值等价性测试。我们建立了一套“黄金测试”Golden Test流程用一个小批量数据在CPU上进行FP32高精度参考计算然后在优化版本上运行逐层、逐张量对比输出确保误差在可接受的范围内如1e-5。坑二内存对齐的“隐形杀手”在实现缓存感知的块分配器时我们最初只保证了起始地址的对齐。但在一次压力测试中当序列长度非常大时性能会出现断崖式下跌。使用硬件性能计数器分析后发现是缓存冲突Cache Thrashing异常严重。原来我们只考虑了单个块的对齐但没有考虑多个块在内存中的相对位置。当大量相同大小的块被分配时它们的地址可能映射到缓存中相同的集合Set导致缓存被频繁换入换出。解决方案我们改进了分配算法在分配时不仅考虑起始对齐还引入了一个伪随机的偏移量或者采用更复杂的地址映射策略使得不同块的地址尽可能均匀地分布在不同的缓存集合中。这个改动让长上下文处理的性能提升了约15%。坑三调度器“优化”引发的活锁为了提升吞吐我们为调度器增加了一个激进策略当有高优先级的短请求到来时会暂停当前长请求的计算先处理短请求。在模拟测试中效果很好。但上线后在真实流量洪峰下系统偶尔会完全卡死。分析发现当持续有短请求涌入时调度器会不停地“抢断”长请求导致长请求永远无法完成形成了“活锁”Livelock。解决方案我们为调度策略增加了“公平性”约束和“饥饿检测”机制。任何一个请求被抢断一定次数后其优先级会被临时强制提升保证它能获得足够的计算资源以完成。这体现了系统设计中的一个核心原则局部最优不等于全局最优调度必须考虑所有请求的完整体验。6. 集成、部署与生态展望将vLLM-Kunlun集成到现有的服务体系中我们采用了容器化的方案。我们构建了专门的Docker镜像其中包含了深度优化的vLLM-Kunlun框架、昆仑芯驱动、KML计算库以及模型服务化接口。部署关键点版本严格锁定昆仑芯的驱动、固件、计算库版本必须严格匹配任何不匹配都可能导致性能下降或运行时错误。我们在CI/CD流水线中设置了严格的版本检查。资源隔离在Kubernetes集群中我们通过设备插件Device Plugin将昆仑芯卡作为可调度资源并为vLLM-Kunlun Pod设置明确的资源请求和限制避免多容器竞争硬件资源。监控与告警我们扩展了Prometheus监控指标除了常规的请求数、延迟、错误率还增加了昆仑芯芯片本身的利用率、温度、内存使用等硬件指标便于提前发现潜在问题。生态展望目前vLLM-Kunlun仍是我们内部的解决方案。但我们看到了其更大的价值。未来我们计划从两个方向推进上游贡献将一些不涉及核心商业机密的、通用性强的优化如部分调度策略改进、通用内存优化技巧尝试贡献给vLLM开源社区推动其更好地支持异构硬件。标准化接口我们正在抽象出一套更清晰的“硬件后端接口”。理想状态下vLLM的核心算法保持不变通过实现不同的后端接口如CUDABackend,KunlunBackend,AscendBackend就能适配不同的硬件。这将大大降低其他国产芯片集成vLLM的难度。这个项目让我深刻体会到在AI基础设施领域软件栈的深度优化与硬件本身同样重要。一颗强大的芯片需要同样强大的软件去“唤醒”它。对于国产算力而言构建一个繁荣、高效、易用的软件生态其紧迫性和重要性丝毫不亚于追求更高的硬件算力峰值。vLLM-Kunlun是我们在这条路上迈出的一小步过程充满挑战但结果令人振奋。它证明了通过深度的、系统性的软件优化国产硬件完全有能力支撑起苛刻的生产级大模型推理负载。这条路还很长需要芯片厂商、框架开发者、应用厂商更紧密的协作。
返回列表