关于阿里云:PAI-BladeLLM推理引擎-超长上下文更高性能

2次阅读

共计 4435 个字符，预计需要花费 12 分钟才能阅读完成。

BladeLLM 是阿里云 PAI 平台提供的大模型推理引擎，致力于让用户轻松部署高性能、低成本的大语言模型服务。BladeLLM 对 LLM 推理和服务的全链路进行了深度的性能优化和工程优化，确保不同模型在不同设施上都达到最优性价比。

除了在惯例上下文长度下的极致性能优化之外，BladeLLM 还冲破了现有 LLM 推理零碎上下文长度的极限，可能反对更长的输出长度以及文本生成长度等，使得 LLM 可能解锁更多的利用场景，并且 BladeLLM 在超长上下文下仍然放弃极致的性能，相比于其余 LLM 推理服务零碎有显著的性能劣势。

本文次要介绍 BladeLLM 在超长上下文方面具备的劣势，包含反对的最大上下文长度以及超长上下文的推理性能。

超长上下文推理能力是 LLM 涌现的重要能力之一，该能力促生了一系列具备微小潜在价值的利用场景，包含个性化的聊天机器人（Character.AI）、文学创作工具（Jasper）、文章摘要工具（ChatPaper）等。个性化的聊天机器人会和用户进行持续性的交互，给予用户工作、情感、学习等多方面的帮忙。LLM 会在交换过程中记忆残缺的聊天内容，模型输出长度逐次递增，在屡次交互后造成超长输出文本序列；文学创作工具借助 LLM 的能力批量生成长篇文本，如小说、故事和剧本等。相比传统的手工创作过程，LLM 文学创作工具能够在短时间内生成大量的背景、情节和对话，在大幅度晋升作家和编剧的创作效率的同时为读者提供更加丰盛且多样的浏览资料。LLM 涌现的超长上下文推理能力被认为是通往 AGI 的必经之路，该能力的意义次要体现在以下几个方面：

摸索更多利用场景：超长文本生成的反对使得 LLM 能够利用于更多的利用场景，如个性化聊天机器人、生成长篇小说、技术文档、学术论文等。这些利用场景通常须要生成较长的文本内容。
生成更具上下文连贯性的文本：LLM 的指标是生成与给定上下文相干的自然语言文本。当生成序列限度较短时，可能会导致生成的文本与上下文的连贯性有余，影响生成文本的品质。而 LLM 反对超长文本生成，能够更好地放弃上下文的完整性，生成的文本更加连贯，从而晋升生成文本的品质。
晋升生成多样性：较长的生成序列能提供更多的空间来摸索不同的文本可能性，从而进步生成文本的多样性。LLM 反对超长文本生成，能够更好地捕获上下文的轻微变动，生成更多样化、丰盛的文本内容。

随着相干利用场景的铺开，反对超长上下文的模型层出不穷，其中包含反对 84K 上下文的 MPT StoryWriter、200K 上下文的 Claude 2 以及 256K 上下文的 LongLLaMA 等等（见下图）。零碎层面尽管曾经有局部框架（如 DeepSpeed）针对超长上下文进行反对和优化，然而仍然集中于训练阶段。而在推理阶段，风行的框架无不面临超长输入输出无奈运行或运行效率低下的问题，能够说超长文本的输出和输入对大模型推理引擎带来新的挑战。

首先，现有的 LLM 推理引擎难以满足大模型解决超长上下文信息的需要，这些零碎对于存储资源的配置计划以及计算算子的设计会极大地限度模型的最大输入输出长度。因而，大规模的上下文反对须要更高效的存储和计算策略；此外，更长的上下文信息使得推理工夫急剧增长，引起成本上升和用户体验的降落，这个问题在现有的 LLM 推理引擎中尤为显著。推理工夫增长的次要起因是 LLM 的 Attention 机制，它须要计算每个 Token 与其余 Token 之间的绝对重要性，随着上下文长度的减少，Attention 计算须要解决更多的 Token 从而导致更长的计算工夫，因而更疾速高效的 Attention 计算方法是减速 LLM 超长文本生成的要害。

以 HuggingFace Llama2-13B 模型为例，随着上下文长度的减少，生成一个 token 的工夫显著减少，具体增长趋势如下图所示。上下文长度 34K 时 HuggingFace 开源模型生成一个 token 的工夫是上下文长度 1K 时的 3.5 倍.

以下是 BladeLLM 推理引擎的技术架构图，蕴含了很多外围组件，本文次要介绍其中的 RaggedAttention 和 DNN-based AutoTuner.

近期，对于 Transformer Multi Head Attention 计算有两个颇具影响力的工作即 FlashAttention 和 PagedAttention, 它们对 LLM 训练和推理零碎的设计范式产生了深远的影响。

PagedAttention 受到操作系统中虚拟内存和分页思维的启发，在不间断的显存空间中存储间断的 keys 和 values. PagedAttention 将每个 sequense 的 kv cache 划分为块，每个块蕴含固定数量的 tokens 的 keys 和 values。因为这些块在显存中不用间断，从而极大地缩小了显存碎片，并且无需为每个 sequense 提前预留大量的显存，使得贵重的显存资源失去了最充沛的利用。极致的显存利用率配合上 Contiguous Batching，极大地晋升了 LLM 推理服务的吞吐。相应地也带来一个毛病，不间断的显存块在肯定水平上影响了 kernel 访存效率，从而影响了性能。

同期 BladeLLM 自研的 RaggedAttention 尽管要解决的问题与 PagedAttention 相似，然而在实现办法上存在肯定差别，具体来说就是在 kernel 性能与显存利用率之间有着不同的 tradeoff。

RaggedAttention 的名字是受 Tensorflow 框架中 RaggedTensor 的启发。Ragged 是不规则的意思，这意味着 RaggedAttention 的 kv cache 不是规定的 Tensor，而是容许其中每个 sequence 的长度各不相同，从而可能和 Contiguous Batching 高效配合，晋升零碎吞吐。然而和 PagedAttention 不同的是，RaggedAttention 保障同一个 sequence 的 key 和 value cache 是间断存储的，因而可能晋升 kernel 的访存效率和进而晋升性能。同样地，间断存储会造成肯定的显存碎片和显存预留问题，从而影响了显存利用率。这是一个典型的工程上的 tradeoff，没有标准答案，因为不同的算力显存配比、不同的输入输出长度、甚至不同业务对于延时的不同要求都会导致系统瓶颈的差别。作为 AI 平台，BladeLLM 致力于为不同模型、不同设施、不同 workload、不同业务场景以自动化的形式寻求最适宜的配置。

例如对于变动范畴极大的上下文长度，借助于下一大节将要介绍的 AutoTuner，RaggedAttention 在不同上下文长度下都能放弃高效的计算和访存，咱们实测上下文长度从 1 变动到 512000，RaggedAttention 都能取得极致的性能。

LLM 推理属于典型的强 Dynamic Shape 场景，不仅 Batch Size 维度会动态变化，Sequence Length 维度变动幅度更为微小。Dynamic Shape 场景下谋求 Kernel 极致性能的次要办法之一是基于理论运行尺寸进行 Tuning 调优，即针对每一组特定的输出尺寸都通过理论运行和测量选取 Best Schedule，采纳这种办法的工作包含 AutoTVM, Ansor 等。这种办法尽管能够达到很极致的性能，然而存在 Tuning 开销大的问题，特地是 Tuning 后果只能对特定 Shape 实用，对于 Dynamic Shape 场景十分不敌对：如果离线事后针对所有可能的 shape 都 tune 一遍，须要破费的 tuning 工夫以及计算资源十分微小；如果在线对每组新 shape 实时进行 tuning，会对线上性能产生重大的性能扰动。

针对以上痛点，BladeLLM 采纳了 DNN-based AutoTuner，齐全依赖 DNN 模型预测的后果而无需理论运行测量来选取 Best Schedule. 咱们在训练数据收集、模型构造、特征提取、Loss 函数设计等方面进行了大量的摸索和尝试，一直晋升 DNN 模型的预测准确率，目前基于 DNN-based AutoTuner 的 GPU 计算密集算子的均匀性能达到基于理论运行测量的 Tuning 调优性能的 99.39%.

在解决了预测准确率之后，升高 DNN 预测模型的运行工夫和占用的计算资源成为该技术利用于高实时性在线推理场景的要害挑战。间接应用已有框架和引擎（如 PyTorch, TorchScript, OnnxRuntime 等）搭建预测模型无奈满足服务的高实时性需要，咱们通过模型零碎联结优化，使得 AutoTuner DNN 模型预测延时升高至 2us. 极致的系统优化使得预测模型性能相比于用 PyTorch, TorchScript, OnnxRuntime 搭建的模型别离晋升 36 倍，19.5 倍和 4.3 倍（见下图），并且推理过程占用的系统资源极低，预测模型只应用一个 CPU Core 而非 GPU 资源以确保不对服务的 GPU 模型本身性能造成任何烦扰。因为微秒级的低预测时延和 99% 以上的预测准确率，AutoTuner 不仅被利用于 LLM 在线推理服务，还胜利服务于包含搜推广、语音辨认、Stable Diffusion 等 Dynamic Shape 业务场景。

咱们以最大文本生成长度以及相应的生成工夫为例来比照不同 LLM 推理零碎最大能反对的上下文长度以及相应的性能，后果如下：

lmDeploy(基于 FasterTransformer)在生成长度超过 10K 之后会 Hang 住
vLLM 在生成长度超过 12K 之后呈现 illegal address 谬误
Huggingface 原始的 Llama 模型在生成长度超过 34K 后 OOM
LightLLM 最大生成长度 (67K) 和 BladeLLM(70K)靠近，然而所须要的工夫是 BladeLLM 的 3 倍

注：为了比照的公平性，以上后果均基于 fp16 权重和 fp16 kv cache 测量，BladeLLM 现已反对 kv cache 量化，可进一步将单卡最大反对的上下文长度晋升至 280K；以上所有的测量均未采纳投机采样；以上测量在 8 月份实现，目前业界 LLM 推理引擎都还在疾速倒退中，咱们期待更新的后果比照，同时 BladeLLM 反对更长上下文、更高性能的新版本开发也靠近序幕，有了新的后果咱们会持续和大家分享。

超长上下文是 LLM 倒退的必然趋势，而以后支流的 LLM 推理和服务引擎所反对的上下文长度以及超长上下文的推理性能都远远不够，以上分享了一些对于 BladeLLM 对超长上下文的反对以及超长上下文推理性能，欢送大家交换探讨。此外，除了关注超长上下文场景，BladeLLM 也会继续关注推理的多个技术方向，包含低比特量化压缩、多轮对话、极致内核优化、编译优化等，后续咱们也会有更多的技术分享对外公开，欢送大家继续关注！

正文完