Modal 发布 Quail：联合优化查询计划与推理引擎加速 AI-SQL

我把它看作 LLM 帕累托最优曲线上的一点，处于一个曾存在大量潜在需求（无需思考、单 token、低延迟、可接受的智能水平）的区间，而由于大家都在竞逐更高智能，这一区间此前投资不足。

- Karpathy 先生，关于 Jev

当每个人和他们的表亲都在大声构建编码智能体和聊天机器人时，后端正在悄然发生一场推理革命。只要成本性能足够好，简单的 LLM 数据转换就能极其强大——只需刷刷社交媒体，看看几个令人瞠目结舌、激发黑客灵感的 TypeSafe AI 的 Jev 模型演示。

Jev 在你可能称之为“JSON 层”的地方实现这些转换，即客户端与服务之间的 Web 风格接口。

AI-SQL 则在分析型 SQL 层实现它，即商业智能与数据库之间的接口：

不同的推理应用会产生不同的推理工作负载，AI-SQL 也不例外。像上面这样的查询可能会产生数百万条包含数千个 token 的序列——你的 token 预算要 RIP 了。这些查询通常所需的智能远低于前沿水平，因此小型开放权重模型就能轻松胜任。但是，天真地将这些序列直接交给一个为智能体推理优化、通过面向任意用户控制请求的接口提供服务的推理引擎，本质上且极其低效。

因此我们构建了一个推理引擎来解决这个问题：查询感知推理层（Quail）。在一个规划尤为重要的多连接查询上，Quail 在单块 H100 GPU 上达到每分钟处理超过十亿 token（TPM/GPU），比同一硬件上的 vLLM 基线快 10 倍以上。在 Modal 上，这相当于每十亿 token 不到 6 美分。

在我们新发布的 AI-SQL 查询基准上，Quail 在各项任务上的几何平均运行速度比 vLLM 快 1.84 倍——其中包括我们为展示 AI-SQL 推理未来改进空间而设计的两个查询。

你现在就可以在 Modal 上试用它：

在这篇博客中，我们将简要概述我们正在解决的问题以及 Quail 目前的工作原理。剧透：最大的优势在于，手握结构化查询，你就可以对请求排序，以更好地缓存（和驱逐）KV。这需要对 Hydragen 风格的级联注意力稍作修改。针对小模型的大量小请求也可能带来大量主机开销，即低GPU 利用率，而当你提前知道请求的结构时，就可以避免这一点。

这是 Modal 的推理研究人员与卡内基梅隆大学全栈数据实验室的数据库研究人员之间的合作——可以称之为“专家混合”。我们分享我们的成果，是因为我们希望让这项工作更加“专家并行”。我们相信，这只是推理与数据库交叉领域开源性能工程的开始——这是计算领域最重要的两个应用。

在本文中，我们将更多地关注推理工程师的考量。你可以从数据库工程师的角度在Full Stack Data Lab 博客上阅读更多内容。你也可以在这里查看 Quail 的代码，或在这里查看文档。如果你大规模运行 AI-SQL 查询，并有兴趣提升性能、降低成本，请联系我们。

什么是 AI Functions 和 AI-SQL？

首先，再补充一些关于该工作负载的背景。

这明确地不是让 AI 系统根据自然语言输入生成 SQL——那是 NL2SQL。那看起来很像传统的聊天机器人或编码智能体工作负载，因此现有的推理引擎可以很好地胜任。

实际上恰恰相反！在 AI-SQL 中，我们使用 SQL 的扩展来以编程方式生成（并消费）AI 系统的提示。提示由数据库条目构建，并生成表。

就像这样：

AI-SQL 主要用于商业智能（BI）平台内部，帮助数据科学家和利益相关者对其半结构化数据（如文档和自由文本字段）提出更“模糊”的问题。

目前还没有标准，但主流托管分析数据库平台都有自己的版本：Snowflake Cortex AI-SQL、Databricks AI Functions、BigQuery AI functions。

与 SQL 的其他部分不同，这个问题实际上需要 GPU。

考虑以下针对 BioDEX 数据集的查询计划，该查询选择对同时包含神经系统和心血管成分的药物产生严重不良事件的报告：

如果这些是普通的过滤和连接，比如基于字符串匹配和逻辑相等，那么即使这是一个分析型查询、看起来“吞吐量很大”，也没有充分理由使用像 GPU 这样的高吞吐量数值加速器。这对某些人来说可能显而易见，但我们还是逐步梳理一下逻辑。

从持久存储加载到内存（或从内存加载到寄存器）的每个字节，最多只需要少量算术/逻辑运算即可实现比较。GPU 是为高算术强度的工作负载设计的——即每加载一个字节要执行许多运算。而最新的 GPU 将大部分算术带宽用于大型矩阵乘法的专用硬件，即Tensor Cores。普通的过滤/连接不需要大型矩阵乘法。

但这个查询计划使用了 AI_FILTER 和 AI_JOIN，它们将输入传递给大型语言模型。LLM 是一系列数值运算，其瓶颈是大型矩阵乘法。从持久存储加载的每个字节在被写入存储之前，将经历数十亿次量级的运算。

为什么这对推理工程师来说很有趣？

如今大多数推理工程都特别聚焦于一种工作负载形态：由用户和推理服务外部的工具调用迭代构建长输入序列。这是聊天机器人和智能体工作负载的形态——也是在对模型进行微调使其成为聊天机器人或智能体的强化学习运行期间“rollout”推理的形态。

别误会，这是非常重要的工作！我们在这里写过我们的方法。但对于硬核推理工程师来说，老实说，它开始让人觉得有点……老生常谈了。

在超低延迟推理方面也有一些工作，其中速度与智能同等重要。我们已在此处撰写了相关技术。一般来说，这些工作负载使用结构化输出/工具调用。它们最终成为推理的“OLTP”，比开放式代理更容易嵌入其他计算机应用。最近 Jev 的流行证明了这些工作负载的重要性——而且我们仍处于早期阶段！

AI-SQL 工作负载尚未受到太多关注——但——我们认为它们对推理工程师来说很有趣，出于一些根本原因，而不仅仅是它们对应用的重要性。最引人入胜的是，它们非常适合 transformers（因为它们能实现“完美”的 KV 缓存使用）以及 GPU 上的 transformers（因为它们不需要解码）。

管理 KV 缓存，无需所有遗憾。

在典型的推理中，请求由客户端控制且任意。这导致了无尽的痛苦。但在 AI-SQL 推理中，客户端只控制 SQL 查询，这些查询会产生许多请求，而组合的查询规划器/推理引擎对这些请求的处理有相当大的控制权。

这使得操作一个能摊销更多工作的缓存变得特别容易。例如，我们确切知道任何缓存条目何时不再需要，因此我们可以无所畏惧地驱逐它。我们还相当了解缓存需求会是什么样子，因为我们一开始就获得了整个查询计划的请求。

我们非常需要为 Transformers 进行缓存，因为它们的前向传递在序列长度上是朴素二次的。通过 KV 缓存，我们可以将其转换为线性时间和线性存储。

对于代理工作负载，KV 缓存可能难以操作，因为访问之间的时间完全未知。但对于 AI-SQL 查询，我们控制推理引擎请求，因此可以预测未来的访问并应用诸如预取之类的优化。此外，由于我们关注令牌吞吐量，我们不太关心检索 KV 条目的延迟。这使得例如操作多级 KV 缓存更加可行。

看，妈妈，没有解码！

序列模型推理分为两个阶段：“预填充”，此时生成大部分 KV 缓存，以及“解码”阶段，此时生成大部分输出令牌。

解码有点麻烦。GPU 并不特别擅长它。解码的算术强度低，因此即使 GPU 提供大量内存带宽，也很难保持算术带宽饱和。

代理应用严重偏向解码——即使输入令牌多于输出令牌，解码速度慢得多，以至于占用了大部分时间。这个问题如此严重，以至于推理服务部署常常被迫采用复杂的解决方案，如跨节点预填充-解码分离，仅仅为了获得可接受的性能。

但并非所有令牌都在解码期间生成。输入序列处理期间的最终“预填充”前向传递会为单个令牌发出预测。

而对于序列的布尔分类，即 AI.IF，单个令牌就是你所需要的——字面意思。

这很重要，因为 AI.IF 不是配角。这就是 AI-SQL 中连接（JOIN ON AI.IF）的实现方式。通过在提示构建中稍加巧妙，AI.CLASSIFY 也可以映射到单个令牌上，对于多达词汇表大小的多个类别（我们已将那个留待未来工作！）。

目前，我们并未充分利用这一点，除了在我们未实现的部分：

• 分离预填充和解码阶段（更不用说解聚合），因为没有解码
• 采样，因为没有生成的令牌，只有概率
• CUDA Graph 捕获，因为预填充的持续时间足够长，启动开销可以忽略不计，即使是在大型 GPU 上运行的小模型
• 推测解码，因为它加速了多个令牌的解码

但我们预期有更深层次的机会来优化仅预填充推理！

Quail 联合优化 SQL 查询和推理工作负载。

掌握了 SQL 问题和推理问题的形态后，让我们快速浏览一下 Quail 的架构，重点关注查询规划器和执行引擎。

架构概览

你可能还没有注意到，但现在构建数据库很容易（参见 Stonebraker & Pavlo, 2024 或 Andrew Lamb 关于 Apache DataFusion 的这次演讲）。具体来说，分析型数据库更容易构建，因为许多关键组件已经通过可扩展的开源实现标准化了。而且，由于编码代理的出现，组合开源组件现在变得极其简单。

关键组件，从外部接口到内部实现细节依次是：SQL 解析器、查询规划器、执行引擎和存储引擎。

1. SQL 解析器。我们使用 @tobymao 的 Python sqlglot 库，它支持 snowflake 和 bigquery 方言。AI-SQL 通过“匿名”表达式处理，即推给查询规划器。
2. 查询规划器。这部分实质上是自定义的，因为它是工作的核心。我们将在下面描述。我们使用 Substrait 来序列化查询计划以进行基准测试。
3. 执行引擎。我们分叉了 vLLM 的模型前向传递实现，然后如下所述修改了内核（主要是编写 Triton 以实现内核融合）。我们还没有添加完整的 SQL 执行支持，但使用 DataFusion 添加是相当直接的。
4. 存储引擎。我们使用 pyarrow 来管理列式 Arrow 格式。这是一个分析型工作负载，即一次写入/多次读取，也就是“简单模式下的文件系统”。我们假设这是从对象存储（如 S3）或分布式文件系统（如 Modal Volumes）预先获取的。

为推理引擎设计查询规划器

查询规划器基于解析 SQL 查询的逻辑计划，并转换该计划——既在等效逻辑计划之间转换，也转换为具有具体操作的“物理”计划。查询规划器的设计空间是巨大的。毕竟，它们本质上是编译器！

但我们的问题集仅限于过滤和连接，在此基础上，我们进一步能够主要使用众所周知的技术。我们进行谓词下推越过连接，基于选择性的过滤排序，如 Hellerstein 和 Stonebraker 所述，以及使用动态规划在深树上进行连接排序，如 1976 年经典的 System R 论文 所述。这是一个非常简短的概述——更多关于数据库方面的细节在 Full Stack Data Lab 博客 这里！

在这里，我们将简要介绍以 Transformer 推理/GPU 为中心的贡献，在成本模型和连接排序算法中。具体来说，Quail 在成本模型中添加了光速估计，并在连接顺序搜索中添加了 KV 感知。

KV 感知的连接顺序搜索

连接顺序传统上基于“分而治之”的动态规划。从高层来看：在某一步选择最优选项，然后在该选择固定的情况下搜索最优子计划。

我们的情况并不像普通的连接顺序那么简单。我们还会额外跟踪之前计划的 KV 状态，以防最初看起来糟糕的选项后来对某个可以复用其 KV 的连接有用。

在搜索过程中，我们维护多个候选方案。只有当新的候选计划具有更少的 token、更少的注意力对和更少的缓存 token 时，我们才会淘汰计划——它在帕累托意义上被“支配”了。然后我们通过应用下文所述的光速成本估算来选择最终计划。

连接顺序很难！我们预计这项技术会有大幅改进，我们很乐意与你一起推进这些改进。

悲观地估算光速

像大多数数据库一样，我们基于硬件的成本估算相当粗糙。我们使用 Williams、Waterman 与 Patterson 的“roofline 模型”来针对面向吞吐的硬件，基于硬件峰值速率估算“光速”，这有其局限性（参见我们 GPU 性能术语表中关于性能瓶颈的条目最后一段）。

但粗糙并不意味着无效！一方面，SoL 模型是迭代期间用于检查结果合理性的关键工具。另一方面，它倾向于高估峰值性能，而不是随机偏差或低估。把它比作北极星：你永远无法到达它，但它仍能帮助你朝北前进。

详细代码在这里，但成本模型大致如下：

该计算仅基于对目标工作负载而言属于关键瓶颈的模块子集：逐 token 注意力投影和 MLP/MoE 层，以及跨 token 注意力计算。

原则上这必须针对每个模型做一次，但在实践中模型共享大量算子。我们发现，当代编码智能体非常擅长读取 Hugging Face 配置，并在这种设置下生成合理的成本模型——尽管它们的输出通常需要凭感觉检查一下。关于基于 roofline-SoL 的成本建模的另一个应用，请参见我们的投机解码加速估算器。

作为执行引擎的推理引擎

一旦选定了最终的物理计划，就必须由执行引擎来实现它。

但在过多考虑 GPU 方面之前，重要的是先让 CPU 不挡路。任何从事过高性能存储或网络工作的人都会熟悉这里的基本节奏。

在这种情况下，我们有大量序列（数百万）和一个大 GPU（H100）上的小模型（数十亿参数），这超出了大多数 tokenizer 后端的设计空间。我们使用了来自斯坦福的我们的朋友 Marcel Rød 的 Gigatoken。附注：这项工作发生在查询规划器中，但随后在执行时被复用。我们将 token ID 存储在内存映射的 Arrow 文件中，这与我们存储引擎中使用的基本技术相同。

深入到 GPU 层：我们从 vLLM 中的模型前向传播实现开始，然后用一些自定义优化重写了它们。我们非常感激能够在此基础上构建社区的工作！

我们用于核心 matmul 和 attention 操作的 kernel 是标准的：DeepSeek 的 DeepGEMM 和 Dao 等人的 Flash Attention 3。这些 kernel 在高吞吐/仅预填充推理方面相当出色。

光速成本模型只考虑这些操作，但前向传播中还有许多其他操作。它们的运行时间原则上可以忽略不计，但数量很多，累积起来也不容忽视，包括启动这些 kernel 的主机开销。

这里的标准技术是将多个较小的 kernel 合并为一个，即kernel 融合。幸运的是，现在用 OpenAI 的 Triton 编写自定义 kernel 相当容易——尤其是在编码代理辅助和参考实现的帮助下！我们将 add-RMSNorm 与 fp8 量化融合，并将每个头的 query/key RMSNorm 与旋转嵌入融合。

我们确实需要将 FlashAttention“包装”在一些自定义逻辑中。在 A 和 B 的连接中，来自 B 的许多后缀文档会关注 A 中的同一个锚点。如果简单地映射到 FlashAttention，在执行过程中会导致锚点的大量重复。

我们通过对连接进行“递归”或“部分组合”的注意力计算来避免这种情况。

首先，我们从所有后缀查询关注锚点，这是一种“交叉注意力”。然后后缀关注自身，最后结果通过 LSE 重新缩放进行合并，类似于 Dao 等人的 flash decoding。后缀的 KV 不会写入缓存——我们进行零解码，并且从不进行 (N>2) 路连接，因此不需要它们！

这与 Flash Infer 的级联注意力 或 Scaly/Hazy 的 Hydragen 非常相似。它看起来也像 树注意力，但针对深度为一的树的特殊情况。我们与这些技术有着相同的动机：共享前缀的重用。但由于我们提前知道共享前缀，我们可以跳过很多复杂性。

此外，我们在模型层面做了一项优化。因为我们只进行过滤和连接，我们只需要模型输出真值和假值 token 的概率。这意味着我们的词汇表只有 8 个选项——这是结构化输出的一个极端案例。这在引擎启动时已知，因此我们直接从模型的语言建模头中删除这些列，将最终的 unembedding matmul 从词汇表大小 x 潜在大小减少到 8 x 潜在大小。

推理引擎的未来是什么？

最后，如果您运行 AI-SQL 查询并有兴趣提高性能和降低成本，请与我们联系。

我们在 SGLang 上做了大量 工作，并构建了自己的自定义引擎，包括 Quail，这让我们对该领域的发展方向有了一些看法。这对推理引擎工作来说是一个关键时期，因为这些系统还很新（只有几年，而不是几十年），对整个软件工程也是如此，因为编码代理正在改变研究和开发的约束条件。以下是一些相关说明。

对于优化的 AI-SQL 推理来说，这仅仅是开始。

对于 Quail 和 AI-SQL 工作负载，还有许多明显的额外优化。我们可以生成更好的计划，而且对于我们执行的计划，我们仍未达到光速。

这里有一个快速的插旗列表，以便我们可以对任何实现它们的人进行 Schmidhuber 我们希望看到更多工作的列表：

1. 为 KV 缓存添加更多层级。

我们只在 GPU HBM 中缓存 KV 值。但还有更多的存储层级，而优秀的缓存总是多层的。像 LMSYS Org 的 HiCache 这样的系统有助于管理多层缓存。我们已将其用于它所设计的工作负载类别，即聊天机器人/智能体推理，但我们在此处并未应用它。

因为这个工作负载相当不同，我们预计现有缓存系统仍有改进空间。特别是，我们对请求顺序既了解得更多，也有更多的控制权，因此我们可以更直接地管理缓存；而且我们对延迟实质上不敏感，因此我们能从更高、更慢的存储层级中受益，尤其是在条带化的情况下。也许我们甚至可以从磁带中为我们的 LLM 提供数据？

2. 跨查询进行优化。

我们也只在请求的生命周期内缓存 KV 值。我们不会把 Bloom 过滤器或 zone map 扔进垃圾桶，那为什么对 KV 缓存要这样做呢？GPU HBM 太宝贵了，不能这样做，但更高的缓存层级要便宜得多。

3. 支持大于内存的数据集。

正如 antirez 的 Redis 和 Mühleisen 与 Raasveldt 的 DuckDB 的成功所证明的那样，你可以在没有持久化后端存储的情况下构建一个有用的数据库系统。而且推理的计算强度如此之大，以至于数据集规模往往更小。

但这不是止步于内存处理的借口！因为 KV 表示比存储表示大得多，应用这项技术也将受益于分层 KV 缓存。

4. 用于更好共享的基数索引。

在我们的基准测试套件中，我们在一个案例上落后于 vLLM：处理智能体轨迹数据集。我们特意添加这个基准，是因为我们想证明我们的系统在最初构建时是在做权衡，而不是以远超现有引擎的工程投入在某种程度上普遍优于它们。

特别是，智能体基准可以映射到智能体服务——只需将所有会话历史存储在数据库中，然后用一条新的输入消息 SELECT 这些历史，并运行 AI_COMPLETE。这会产生大量跨文档前缀共享，而我们当前的系统无法建模或利用这一点。

但我们仍然可以让 Quail 在这里的性能更好！例如，我们可以在文档上构建基数树索引，然后寻找 KV 缓存复用的机会。

5. 更好的跨内核重叠。

我们坚持使用相对简单的内核级技术，比如 Triton 中的融合，并且我们知道我们仍然远未达到光速。造成这种差距的一个常见原因是跨内核的开销或资源复用不足。一个面向吞吐量的、类似 Hazy Research 的 megakernel，甚至只是 CUDA 编程式依赖启动，都会改善重叠，并可能提升性能。

6. 即时微调。

最后，是一些更具推测性的内容。在现代分析数据库中，算子常常在执行中途被优化，例如通过特化和 JIT 编译——参见 Umbra 的“Flying Start”编译器。这些优化后的模块随后会在查询执行中途被换入。

我们可能同样通过量化、剪枝、蒸馏及类似技术“编译”出更快的模型，与查询并发运行，然后在准确率达到可接受水平后换入。而 ML 训练确实可以被操作化到这种程度！参见 Anil 等人的“On the Factory Floor”论文。

以共享核心进行 vibe 编码的推理引擎。

我们大部分工作是通过提示编码代理完成的。难点不再是执行想法，而是问题定义和结果度量——品味和质量保证。

软件工程和系统研究正朝着这个方向发展吗？我们不知道，但对我们来说似乎很有可能！无论发生什么，都会很奇怪，并且需要对工作习惯、社区结构和激励机制进行深思熟虑的关注，这一点最近已被数学界令人钦佩地展示出来。

更狭义地说，我们非常清楚，最近定制推理引擎如雨后春笋般涌现，这些引擎是针对更窄的工作负载或工作负载集而开发的。SyFI 实验室的 VibeServe 代理开发系统表明，这种创作过程甚至可以泛化，前提是正确定义基线和目标。

我们预计这一趋势将持续下去，但我们还不指望所有推理系统代码都针对每个应用从头编写。开源协作带来的好处太多了——运维简单、速度提升、“只要眼球足够多，所有 bug 都浅显易见”。编码代理改变了一些系数，但我们不认为它们会消除驱动开源贡献的合作均衡策略。

相反，我们认为推理世界很快就会更像 DataFusion 之后的数据库：一个可复用的“核心”，既足够表达力强以吸收新技术，又足够受控以提供保证。作为概念验证，请参见 ekzhang 关于代理可扩展推理引擎的推文。

像 vLLM 和 SGLang 这样的现有推理引擎也可能被改造，以更好地充当那个核心——参见 nano-vllm 和 mini-sglang 项目。

致谢

我们感谢 Adobe 的 Joe Barrow 以及 DoubleWord 的团队，特别是 Fergus Finn，审阅并提供了对本工作的反馈。