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

除了在常规上下文长度下的极致性能优化之外，BladeLLM 还突破了现有 LLM 推理系统上下文长度的极限，能够支持更长的输入长度以及文本生成长度等，使得 LLM 能够解锁更多的应用场景，并且 BladeLLM 在超长上下文下依然保持极致的性能，相比于其他 LLM 推理服务系统有显著的性能优势。

本文主要介绍 BladeLLM 在超长上下文方面具有的优势，包括支持的最大上下文长度以及超长上下文的推理性能。

背景

超长上下文是 LLM 发展的必然趋势

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

1. 探索更多应用场景：超长文本生成的支持使得 LLM 可以应用于更多的应用场景，如个性化聊天机器人、生成长篇小说、技术文档、学术论文等。这些应用场景通常需要生成较长的文本内容。

2. 生成更具上下文连贯性的文本：LLM 的目标是生成与给定上下文相关的自然语言文本。当生成序列限制较短时，可能会导致生成的文本与上下文的连贯性不足，影响生成文本的质量。而 LLM 支持超长文本生成，可以更好地保持上下文的完整性，生成的文本更加连贯，从而提升生成文本的质量。

3. 提升生成多样性：较长的生成序列能提供更多的空间来探索不同的文本可能性，从而提高生成文本的多样性。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.

RaggedAttention

近期，关于 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 都能获得极致的性能。

DNN-based AutoTuner

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 也会持续关注推理的多个技术方向，包括低比特量化压缩、多轮对话、极致内核优化、编译优化等，后续我们也会有更多的技术分享对外公开，欢迎大家持续关注！