通过结合查询规划器和推理引擎，在单个GPU上实现每分钟十亿令牌的处理速度（15分钟阅读时间）

# 通过组合查询规划器与推理引擎，在单个GPU上实现每分钟十亿tokens处理量 | Modal博客 来源：https://modal.com/blog/quail-billion-tpm > *我将其视为LLM帕累托最优曲线上的一个特定点，该点对应着存在巨大潜在需求（无需思考、单token输出、可接受低延迟智能）的领域，但由于追求更高智能的竞赛，这一领域投资不足。* - 卡帕西谈Jev (https://x.com/karpathy/status/2102124533729955960?s=20) 当众人及其亲友都在大张旗鼓地构建编程代理和聊天机器人时，一场更为静默的推理革命正在后端进行。只要性价比足够高，简单的LLM数据转换就能发挥惊人威力——只需浏览社交媒体，就能看到一些令人瞠目结舌、启发黑客的TypeSafe AI (https://typesafe.ai/) Jev (https://typesafe.ai/blog/introducing-system-one-models-and-jev)模型演示。 Jev在所谓的“JSON层”实现这些转换，即客户端与服务间的Web风格接口。 AI-SQL (https://docs.snowflake.com/en/user-guide/snowflake-cortex/aisql)则在分析SQL层——商业智能与数据库的接口层实现这一点： 不同的推理应用会产生不同的推理工作负载 (https://modal.com/llm-almanac/workloads)，AI-SQL也不例外。类似上述的查询可能生成数百万个数千tokens的序列——足以耗尽您的token预算。这些查询通常不需要前沿级智能，因此小型开放权重模型就能出色完成任务。但天真地将这些序列直接通过面向任意用户控制请求的接口，传递给为代理推理优化的推理引擎，本质上是极其低效的。 因此我们构建了一个推理引擎来解决这个问题：查询感知推理层 (https://github.com/fsdatalab/quail)（Quail）。在一个规划尤为重要的多连接查询中，Quail在单个H100 GPU上达到了每分钟超过十亿tokens的处理量（TPM/GPU），比我们在相同硬件上的vLLM基线快10倍以上。在Modal平台上，这意味着每十亿tokens的成本低于6美分。 在我们新发布的AI-SQL查询基准测试 (https://github.com/fsdatalab/quail-bench)中，Quail比vLLM快1.84倍（任务几何平均）——包括我们设计的两个旨在展示AI-SQL推理未来改进方向的查询。 您现在就可以在Modal上试用它： 在本篇博客中，我们将快速概述我们正在解决的问题以及Quail当前的工作方式。剧透：最大的优势在于，当拥有结构化查询时，您可以对请求进行排序以更好地缓存（并淘汰）KV缓存。这需要对Hydragen (https://arxiv.org/abs/2402.05099)风格的级联注意力 (https://flashinfer.ai/2024/02/02/cascade-inference.html)进行轻微修改。大量针对小模型的小型请求也会产生大量主机开销 (https://modal.com/blog/host-overhead-inference-efficiency)，即导致GPU利用率 (https://modal.com/blog/gpu-utilization-guide)低下，而当您提前知晓请求结构时，就可以避免这一点。 这是Modal的推理研究员与卡内基梅隆大学全栈数据实验室 (https://fsdatalab.github.io/)的数据库研究员之间的合作——称之为“专家混合”。我们分享所做的工作，是希望使其更“专家并行”。我们相信，这只是推理与数据库交叉领域开源性能工程的开端——这是计算领域最重要的两个应用。 本文将更侧重于推理工程师的考量。您可以从数据库工程师的角度阅读更多内容，见全栈数据实验室博客 (https://fsdatalab.github.io/blog/introducing-quail/#43-quail-dominates-vllm-on-bio-4-1404x-faster)。您也可以在此 (https://github.com/fsdatalab/quail)查看Quail的代码，或在此 (https://fsdatalab.github.io/quail/docs)查看文档。如果您大规模运行AI-SQL查询并有兴趣提高性能和降低成本，请联系我们 (https://modal.com/blog/quail-billion-tpm#)。 ## 什么是AI函数和AI-SQL？ 首先，进一步介绍工作负载背景。 这**绝对不是**提示AI系统根据自然语言输入生成SQL——那是NL2SQL (https://arxiv.org/html/2408.05109v4)。这看起来很像传统的聊天机器人或编程代理工作负载，因此现有的推理引擎表现良好。 实际上恰恰相反！在AI-SQL中，我们使用SQL的扩展来程序化地生成（并消费）给AI系统的提示。提示从数据库条目构建并生成表格。 例如： AI-SQL主要用于商业智能平台内部，帮助数据科学家和利益相关者对其半结构化数据（如文档和自由文本字段）提出更“模糊”的问题。 目前还没有标准（但），主要的托管分析数据库平台有各自的风格：Snowflake Cortex AI-SQL (https://docs.snowflake.com/en/user-guide/snowflake-cortex/aisql)、Databricks AI Functions (https://docs.databricks.com/aws/en/large-language-models/ai-functions)、BigQuery AI functions (https://cloud.google.com/blog/products/data-analytics/sql-reimagined-for-the-ai-era-with-bigquery-ai-functions)。 ## 与SQL的其他部分不同，此问题实际上需要GPU。 考虑针对BioDEX数据集 (https://github.com/KarelDO/BioDEX)查询的以下计划，该计划选择包含神经系统和心血管系统组成部分的药物严重不良事件报告： 如果这些是基于字符串匹配和逻辑等价的普通过滤器和连接，则没有充分理由使用像GPU这样的高吞吐量数值加速器，即使这是分析查询（看起来“吞吐量高”）。这对某些人可能是显而易见的，但我们还是逐步分析一下逻辑。 从持久存储加载到内存的每个字节（或从内存到寄存器）最多只需要少量算术/逻辑操作来实现比较。GPU设计用于具有高算术强度 (https://modal.com/gpu-glossary/perf/arithmetic-intensity)的工作负载——即每个加载字节执行大量操作。最新的GPU的大部分算术带宽 (https://modal.com/gpu-glossary/perf/arithmetic-bandwidth)都位于用于大规模矩阵乘法的专用硬件中，即Tensor Cores (https://modal.com/gpu-glossary/device-hardware/tensor-core)。普通的过滤/连接不需要大规模矩阵乘法。 但此查询计划使用`AI_FILTER`和`AI_JOIN`，它们将输入通过大型语言模型处理。LLM是一系列数值操作，其瓶颈在于大规模矩阵乘法。从持久存储加载的每个字节在写入存储之前，将经历数十亿次量级的操作。 ## 这为何对推理工程师有趣？ 如今大多数推理工程都专注于一种特定的工作负载模式：通过推理服务外部的用户和工具调用迭代构成长输入序列。这是聊天机器人和代理的工作负载模式——也是在微调模型成为聊天机器人或代理的强化学习运行期间的“展开”推理。 别误会，这是非常重要的工作！我们在此 (https://modal.com/blog/trillion-tokens-trillion-parameters)介绍了我们的方法。但对于硬核推理工程师来说，老实说，这开始感觉有点……老生常谈。 还有一些关于超低延迟推理的工作，其中速度与智能同等重要。我们在此 (https://modal.com/blog/achieve-sota-specdec)介绍了我们的技术。通常，这些工作负载使用结构化输出/工具调用。它们最终类似于推理的“OLTP”，比开放式代理更容易嵌入其他计算机应用。Jev最近的流行展示了这些工作负载的重要性——并且我们仍处于非常早期的阶段。 AI-SQL工作负载尚未获得如此多的关注——但我们认为它们对推理工程师有几个根本原因而有趣，这与它们对应用的重要性截然不同。最引人入胜的是，它们非常适合Transformer（因为它们能实现“完美”的KV缓存使用）以及GPU上的Transformer（因为它们不需要解码）。 ## 无后顾之忧地管理KV缓存。 在典型推理中，请求由客户端控制且任意。这导致了无尽的痛苦 (https://modal.com/blog/trillion-tokens-trillion-parameters)。但在AI-SQL推理中，客户端仅控制SQL查询，这些查询创建许多请求，而组合的查询规划器/推理引擎对这些请求的处理拥有相当大的控制权。 这使得操作缓存以摊销更多工作变得特别容易。例如，我们确切知道何时任何缓存条目不再需要，因此可以放心地淘汰它。我们也相当清楚缓存需求的样子，因为我们预先获得了整个查询计划的请求。 而我们迫切需要为Transformer进行缓存，因为它们的向前传播在序列长度上天然是二次方复杂度。我们可以通过KV缓存将其转换为线性时间和线性存储。 对于代理工作负载，KV缓存可能难以操作，因为访问之间的时间完全未知。但对于AI-SQL查询，我们控制推理引擎请求，因此可以预测未来的访问并应用预取等优化。此外，由于我们关注的是token吞吐量，我们不太关心检索KV条目的延迟。这使得例如操作多级KV缓存变得更加可行。 ## 看，妈妈，没有解码！ 序列模型推断分为两个阶段：“预填充”，此时生成大部分KV缓存；以及“解码”阶段，此时生成大部分输出token。 解码有点麻烦。GPU并不特别擅长解码。解码的算术强度 (https://modal.com/gpu-glossary/perf/arithmetic-intensity)很低，因此即使GPU提供了大量内存带宽 (https://modal.com/gpu-glossary/perf/memory-bandwidth)，也很难保持算术带宽 (https://modal.com/gpu-gpu-glossary/perf/arithmetic-bandwidth)饱和。 代理应用偏向解码阶段——尽管输入token比输出token多，但解码慢得多，因此占据了大部分时间。这个问题如此严重，以至于推理服务部署常常被迫采用复杂的解决方案，如跨节点预填充-解码分离，以获得可接受的性能。 但并非*所有*token都在解码期间生成。输入序列处理期间的最终“预填充”向前传播会生成单个token的预测。 而对于序列的布尔分类，即`AI.IF`，您只需要一个token——字面上如此。 这很重要，因为`AI.IF`并非配角。它是AI-SQL中连接的实现方式（`JOIN ON AI.IF`）。通过在提示构建上稍加巧思，`AI.CLASSIFY`也可以映射到单个token，类别数量可达词汇表大小（我们将这个留作未来工作！）。 目前，我们并没有充分利用这一点，除了在我们*未*实现的方面： - 分离的预填充和解码阶段（更不用说分离），因为根本没有解码 - 采样，因为没有生成的token，只有概率 - CUDA Graph捕获，因为预填充持续时间足够长，即使在大型GPU上的小模型，启动开销也可以忽略不计 - 推测解码 (https://modal.com/blog/spec-is-all-u-need)，因为那加速的是多于一个token的解码 但我们期待深入优化仅预填充推理的机会！ ## Quail联合优化SQL查询和推理工作负载。 在掌握了SQL问题和推理问题的模式后，现在让我们快速介绍Quail的架构，重点放在查询规划器和执行引擎上。 ## 架构概述 您可能还没注意到，现在构建数据库很容易了（参见Stonebraker & Pavlo, 2024 (https://dl.acm.org/doi/10.1145/3685980.3685984) 或Andrew Lamb关于Apache DataFusion的演讲 (https://www.youtube.com/watch?v=iJhRbDFJjbg)）。具体来说，分析数据库更容易构建，因为许多关键组件已标准化，并有可扩展的开源实现。而得益于编程代理，组合开源组件现在变得容易。 关键组件按从外部接口到内部实现细节的顺序是：SQL解析器、查询规划器、执行引擎和存储引擎。 1. **SQL解析器**。我们使用`@tobymao` (https://github.com/tobymao/sqlglot)的Python `sqlglot`库，它支持`snowflake`和`bigquery`方言。AI-SQL通过“匿名”表达式处理，即转交给查询规划器。 2. **查询规划器**。这部分是实质性的定制，因为它是工作的核心。我们在下面描述它。我们使用Substrait (https://substrait.io/)序列化查询计划用于基准测试。 3. **执行引擎**。我们从vLLM的模型向前传播实现中分叉出来，然后如下所述修改内核（主要编写Triton以实现内核融合）。我们尚未添加完整的SQL执行支持，但通过DataFusion可以相当直接地添加。 4. **存储引擎**。我们使用`pyarrow` (https://github.com/apache/arrow)管理列式Arrow格式。这是分析工作负载，即一次写入多次读取，也就是“简单模式的文件系统”。我们假设这预先从对象存储（如S3）或分布式文件系统（如Modal Volumes (https://modal.com/docs/guide/volumes)）获取。 ## 为推理引擎设计查询规划器 查询规划器获取基于SQL查询解析的逻辑计划，并转换该计划——跨越等效逻辑计划并转换为包含具体操作的“物理”计划。查询规划器的设计空间巨大。毕竟，它们本质上是编译器！ 但我们的问题集仅限于过滤器和连接，并且在其中我们进一步能够主要使用众所周知的技术。我们进行谓词下推穿过连接，基于选择性进行过滤排序（如Hellerstein和Stonebraker (https://dsf.berkeley.edu/jmh/miscpapers/sigmod93.pdf)），以及使用动态规划进行连接排序（如经典1976年System R论文 (https://dl.acm.org/doi/10.145/320455.320457) 中的深树）。这是一个非常简略的概述——更多数据库方面的细节见全栈数据实验室博客 (https://fsdatalab.github.io/blog/introducing-quail/)！ 在此，我们将简要提及Transformer推理/GPU为中心的贡献，包括成本模型和连接排序算法。具体来说，Quail在成本模型中增加了光速估计，在连接排序搜索中增加了KV感知。 ### KV感知的连接排序搜索 连接排序传统上基于“分而治之”的动态规划。高层来说：在一步中选择最优选择，然后搜索固定该选择的最优子计划。 我们的情况不完全像普通的连接排序那样简单。我们还跟踪之前计划的KV状态，以防最初看起来不好的选项最终对可以重用其KV的后续连接有用。 在搜索过程中，我们维护多个候选方案。我们仅在新候选计划具有更少token、更少注意力...