SLM 生成式高并发推理引擎:延迟约束下的有效吞吐优化

SLM 生成式高并发推理引擎:延迟约束下的有效吞吐优化


搜索和推荐场景中想要使用语言模型提高在线效果,往往面临着在线极低延迟的约束,现有通用大模型推理引擎并非在这一延迟约束下设计,因此本文从「具有确定 decode 次数上界的请求」这一工作负载出发,重新设计在线执行路径:以静态 DAG 组织完整业务请求,结合 AOTInductor 编译、Packed Prefill、请求内 K/V 共享及选择性 FP8,系统性压缩调度与计算开销。在平均端到端时延低于 20ms 的约束下,Prefill-only 有效 MFU 达到 67%,吞吐较 SGLang 提升 10 倍,为 SLM 在搜推场景的大规模落地提供了一套高性能、可复用的工程实践。

摘要

搜推场景下业界逐渐使用 SLM(Small Language Model, 小参数量语言模型0.5~1B),用于意图理解,相关性判断,生成式推荐等任务,这些模型保留 Decoder-only Transformer 和 Next-Token Prediction 的基本结构,但在线执行形态与开放式文本生成存在明显差异:

  • Prefill-Only 模式:仅执行一次 prefill,根据每条序列最后一个有效 Token 位置的 logits 得到类别或分值,或基于 prefill 产出的 Embedding 接入 MLP 模块做分类任务;
  • Beam Search 模式:执行一次 prefill 和固定次数 decode,在固定 beam 宽度下返回少量候选序列。两类请求的输入规模、运行状态和最大迭代次数均具有可预先确定的上界。

本文将这类控制流结构、状态容量和执行资源能够在部署前界定,输出 decode 的次数具有固定上界的请求,统称为有界生成负载;其中Prefill-Only 可以看作为 decode 次数为 0 的特殊情况;

在项目早期,我们通过 Triton Inference Server 的 Python Backend 适配业务协议,并使用 SGLang 完成模型执行。该方案能够快速复用 tokenizer、生成调度和状态管理能力,支持业务快速实验拿到效果。但随着 SLM 在搜索场景的持续应用,越来越多的场景取得显著收益,推全成本也在显著增加,业界现有方案已经不足以支持更大规模的落地 SLM,我们随即基于模型负载的独特性,提出专为有界生成负载设计的推理系统。

我们构建了一套以静态有向无环图(Directed Acyclic Graph, DAG)为在线执行抽象的推理框架,将完整业务请求作为 DAG 执行、端到端测量和容量规划的共同边界,并以节点作为模型编译边界。模型计算节点使用 PyTorch AOTInductor 生成的 AOTI 产物,文本处理和数据转换由框架自定义节点完成,Runtime 负责节点调度、资源管理和并发执行。针对有界生成负载,系统进一步将 Packed Prefill、请求内 shared-prefix K/V 布局、固定 Beam Search 编译和选择性 FP8 矩阵计算等自定义实现嵌入 AOTInductor 编译流程。

性能结论表明,约 0.6B 的 Prefill-Only 模型在每个请求包含 15 条样本时达到约 100 request QPS,即约 1,500 sample/s;约 0.5B 的短 prompt 和长 prompt Beam Search 模型在 beam width 约为 100、decode 4 次时,分别达到约 300 和 120 request QPS。三类负载的平均端到端时延均低于 20 ms,P99 约为 70 ms,Prefill-Only 负载下混合精度有效 MFU 为 67%,约为 SGLang 10 倍,短 prompt beam search 有效MFU 41%,约为 SGLang 6 倍。

目前,该框架已经在美团搜索意图分流、搜索相关性判断、生成式召回等应用 SLM 的场景取得实验显著性,能支撑业务在推理成本可控的情况下全量推送。

01 有界生成负载的特征与设计约束

本文关注搜索推荐场景下一类在线执行过程受到业务约束的 SLM 负载:模型仍采用 Decoder-only Transformer,但一次请求涉及的模型阶段和最大迭代次数固定,搜索宽度和状态容量也具有可在部署前确定的上界。实际输入长度、batch 大小和搜索候选仍能随请求变化,但这些动态性均发生在预先配置的容量范围内。

“执行结构固定、动态 shape 可控”是后续系统设计的基础。这些约束来自具体业务的在线执行方式,与模型参数量或模型结构没有必然关系;同一个语言模型用于开放式文本生成时,也可能不满足这里的有界条件。

1.1 两类在线执行形式

Prefill-Only 模型只执行一次 Transformer prefill,不进入自回归 decode。模型在每条序列最后一个有效 Token 位置计算 Next-Token Logits,再根据目标标签对应 Token 的 logits,通过 softmax、Top-K 或固定映射得到业务输出。该路径仍使用语言模型的 Next-Token Prediction Head,但在线执行只包含一次批量前向过程。Prefill 结束后即可返回结果,不需要跨阶段维护该请求的 K/V 状态。

Beam Search 模型在 prompt prefill 后继续执行固定次数 decode。首步 logits 用于初始化候选,后续每一步从各 beam 的 Token 分布中选择局部候选,与历史累计分数合并,再保留固定数量的候选序列。对于一个具体部署配置,beam 宽度和最大生成 Token 数均为固定值。Beam Search 中的候选、累计分数和 parent index 仍由模型输出决定,但状态容量、更新次数和算子组合可以提前确定。

SLM 生成式高并发推理引擎:延迟约束下的有效吞吐优化

1.2 输入与状态边界

本文覆盖的请求平均输入长度均在 500 Token 以内,输出则是分类分值或少量生成 Token。部分 Prefill-Only 请求包含多条输入样本;有界 Beam Search 的 beam 宽度和最大生成步数由部署配置确定,本文覆盖的配置不超过 400 beam 和 5 个生成 Token。两类请求的实际输入和候选内容会变化,但处理规模和状态容量都具有明确上界。

对于同一 batch 内的变长输入,统一 padding 会使短序列承担与最长序列相同的线性层、MLP 和 attention 计算。系统因此保留每条序列的有效长度,并在后续 Prefill 计算中只处理有效 Token。

Prefill-Only 请求在一次前向后结束,不需要跨阶段保存 K/V。有界 Beam Search 只在少量 decode step 之间维护状态,同一请求内的 beam 共用一份 prompt K/V,并分别保存较短的 suffix K/V。输入长度、beam 宽度和生成步数均有上界,因此所需状态可以在加载阶段完成容量规划和内存预分配。

跨请求 prefix cache 的收益取决于公共 Token 前缀的长度和命中率。当前模板较早进入业务 Query、候选文本等请求特有内容,不同请求间可复用的公共前缀很短,实际命中率也很低;同时,平均输入长度不足 500 Token,缓存管理能够摊薄的计算有限。因此,系统不维护跨请求 Prefix Cache,只利用同一 Beam Search 请求内部确定存在的 prompt K/V 共享。

1.3 在线工作负载分析

由于 GPU 推理本质上是基于 Cuda Stream 的串行计算,因此推理服务的耗时可以抽象为排队论中的 G/G/1 模型。排队等待时间的经典工具是 Kingman 近似公式(G/G/1 排队): E[Wq]≈ρ1−ρ⋅ca2+cs22⋅τ\mathbb{E}[W_q] \approx \frac{\rho}{1 – \rho} \cdot \frac{c_a^2 + c_s^2}{2} \cdot \tau, 线上统计 QPS 是每秒处理的请求数,因此衡量推理性能的时间窗口就是 1000ms, 则可以对公式进行如下解释:

  • ρ\rho 可以表征 GPU 的拥挤程度,表示在 1000ms 内执行计算的时间占比。
  • cac_a 请求到达间隔的变异系数,可以表征的是在 1000ms 内请求到达 GPU的概率分布,通常自然流量遵循泊松分布,但由于序列过长而需要分批处理时则会偏离泊松分布,请求在 1000ms 内的分布出现聚集现象但 GPU 处理能力有限,造成排队。
  • csc_s 服务时间变异系数, 表征的是对于推理服务处理的耗时分布情况, 如果序列长度,batch size 等分布不均,则会使得单次请求处理时间分布也呈现出方差过大的现象。
  • τ\tau 表征单次 forward 的平均耗时,表示单次请求占用 GPU 卡的平均时间,可以表示为请求处理的周转速度。

推理耗时可以分为四部分,CPU 入口等待时间、CPU 处理时间、GPU 入口等待时间、GPU 处理时间。 CPU 通常并发处理能力远强于 GPU,因此 GPU 成为资源瓶颈。 越大, GPU 入口等待时间就会延长,当超过上游超时后,排队中的请求被 GPU 顺序执行则是一种浪费,大量请求的丢弃又会导致 GPU 在 1000ms 内的处理时间不足,也就是 ρ\rho 偏低。于是呈现出一种 GPU 利用率很低但耗时却很高的现象。

在搜推场景下,由于序列分批导致请求到达间隔严重偏离泊松分布,服务时间通常是稳定的(序列长度和 batchsize 基本固定),单次推理的平均耗时一般在 20ms 以下。 传统的 SGLang/vLLM框架并不是在这样的假设下设计,调度策略难以处理严重偏离泊松分布的请求到达负载,而其为了prefill 长度任意变化,decode 次数无上界等场景设计的各种优化策略,在搜推的 SLM 场景均是负向优化。

我们利用 csc_s 在搜推场景下较小的优势,移除 SGLang 为无界生成负载定制的各种优化手段,重新为有界生成负载(Prefill-Only & Beam Search)设计 GPU 计算路径,以达到减少 τ\tau 的目的,通过移除 Python 对 request 和 Token 的调度开销压缩 CPU 等待与处理时间,从而达到减少 GPU 空泡的目的,进一步通过 cpp 精细化的 Multi-Stream 调度能力,增加 GPU 跨请求并行计算的能力,以缓解无法调和的 cac_a。

1.4 负载特征与系统设计

上述业务约束将系统问题收束为:如何高效处理边界已知的一批输入。编译、执行和容量控制均围绕这一目标组织。

02 基于静态 DAG 的分层编译执行架构

整体上,本文实现的是一套以静态 DAG 组织完整请求流程、以独立节点封装具体计算的分层编译执行架构。离线构建链使用 Python 表达模型语义、选择算子实现,并将完整 GPU 计算编译为可直接加载的 AOTI 节点产物;在线执行链使用 C++ 完成协议处理、DAG 调度、资源管理和数据转换,不再解释或即时构建模型计算。

在线执行链进一步分为服务、DAG Runtime 和节点三个层次。服务层负责 RPC 协议、请求准入和生命周期;DAG Runtime 根据静态拓扑调度节点,并管理跨节点 Tensor、设备同步、请求上下文和并发;节点负责具体计算,其中 CPU 节点承载文本处理和结果编码,GPU 节点加载 AOTI 编译产物。Runtime 只依赖节点的输入输出、设备和执行接口,attention、Beam Search 和精度策略都封装在节点内部。

离线构建:Torch 模型与算子语义
            -> lowering 与 AOTInductor
            -> AOTI 模型节点产物

          C++ 节点实现与配置
            -> CPU 自定义节点与资源

在线执行:RPC 服务
            -> DAG Runtime
                 -> CPU 预处理节点
                 -> AOTI 模型节点
                 -> CPU 结果节点
            -> RPC 返回

两类模型在这套架构中对应两张不同的 DAG:

这种分层首先使系统可以按完整业务请求执行:Prefill-Only 的多条样本一次进入模型节点,Beam Search 的各个 decode step 在生成节点内部完成。其次,静态拓扑和有界 shape 使 buffer、stream 和执行实例能够在加载阶段规划,请求热路径只处理有效数据。最后,模型计算被封装在节点内部,算子、精度和硬件实现可以独立替换,而服务协议和 DAG 调度机制保持稳定。

2.1 完整业务请求的执行边界

项目早期通过 Triton Inference Server 的 Python Backend 适配业务协议,并使用 SGLang 完成模型执行。这套方案能够快速复用 tokenizer、生成调度和状态管理能力,适合模型验证和早期上线。但在当时的接入方式中,批量输入会被拆分为多条独立生成请求,分别完成文本处理、状态创建和模型调度,返回后再恢复原有顺序和业务结构;Beam Search 模型的各个 decode step 也继续经过通用生成控制路径。

通用生成系统通常按照独立请求组织调度,并在运行时处理请求的加入、退出和解码长度变化。这种粒度适合开放式生成;对于当前只执行一次 prefill 或少量 decode 的负载,请求拆分、对象创建、进程通信、逐步控制和结果重组却都要在约 20 ms 的端到端时延预算内完成,较长生成任务中能够被模型计算摊薄的固定开销在这里更加显著。

因此我们重新选择执行单位:完整业务请求成为 DAG 调度、端到端测量和线上容量规划的共同边界。Prefill-Only 请求内部可以包含多条样本,但这些样本作为一个 batch 进入模型节点,不再拆成独立生成任务;Beam Search 请求则从 prompt prefill 到最终候选返回只经历一次 DAG 调度。模型编译仍以节点为边界,不与请求调度耦合。

项目早期的执行路径可以概括为:

批量业务输入
  -> 拆分为逐样本生成请求
  -> tokenizer / 请求状态创建
  -> 通用生成 Runtime
  -> 逐 decode step 控制
  -> 结果重组

当前执行路径则为:

业务请求
  -> 协议适配与请求准入
  -> DAG Runtime
       -> 文本处理与 packed tokenizer
       -> prefill / 生成 AOTI 节点
       -> 结果编码
  -> RPC 输出

Prefill-Only 中的多条序列由模型节点一次处理,Beam Search 的逐步控制则收进生成节点。在线请求不再反复跨越 Python Backend、服务进程和通用生成 Runtime,协议处理、DAG 调度、buffer 管理及 CPU 数据转换均在进程内完成。这一执行路径减少了请求拆分、对象创建、进程通信、逐步调度和结果重组产生的固定开销。

2.2 静态 DAG 的节点执行模型

静态 DAG 描述协议适配后的节点拓扑、跨节点 Tensor、执行设备和同步关系,而不要求每次请求具有完全相同的数据规模。Runtime 按照以下层次组织执行:

DAG Runtime
  -> Node
       -> CPU Engine  -> 节点配置与本地资源
       -> AOTI Engine -> 编译模型产物

DAG 声明节点之间的依赖和 Tensor 接口,节点类型决定使用哪类执行引擎,执行引擎负责调用具体实现。AOTI Engine 加载模型编译产物,CPU Engine 则调用 tokenizer、数据转换和结果编码等自定义节点,并加载各自的配置与资源。

公开协议中的字符串不会直接作为模型节点输入。在线服务首先将 String 输入转换为连续 UTF-8 payload 和 lengths,DAG 内部再通过 CPU 节点完成规范化、tokenizer 和 packed 数据组织;DAG 输出到 RPC response 的最终映射同样由服务接口完成。由此,DAG 的输入输出保持明确的 Tensor 表示,同时不将具体 RPC 协议耦合到节点实现中。

静态 DAG 固定的是执行结构和物理资源边界,并不要求 Tensor shape 完全固定。系统在加载阶段确定:

  • 节点拓扑和依赖关系;
  • 跨节点 Tensor 的类型与接口;
  • Tensor 的最大物理容量;
  • 最大迭代结构;
  • 执行设备和同步边界。

请求阶段仍然动态确定:

  • 请求内实际 batch 大小;
  • 每条文本的有效 Token 数;
  • 本次 Tensor 的逻辑 shape;
  • Beam Search 中由模型输出决定的状态内容。

这一执行模型以固定结构承载有界的动态数据。加载阶段根据动态 shape 上界和执行并发预分配物理资源;请求到达后,CPU 节点只写入本次有效数据并更新逻辑 shape,下游节点仍按实际数据范围执行。输入长度和 Beam 状态等业务动态性得以保留,同时不会改变节点拓扑和物理资源容量。

节点边界主要由跨节点数据表示、执行设备和同步关系决定。字符串规范化、tokenizer 和结果编码由 CPU 节点执行。Prefill-Only 的 AOTI 节点包含 Transformer prefill 和分值计算;Beam Search 的 AOTI 节点则包含 Transformer、候选选择和固定 Beam 状态更新。一个 AOTI 节点可以包含多个算子及其固定组合,DAG Runtime 只负责调用节点,不需要理解内部的模型拓扑、搜索状态或精度策略。

同一请求上下文中的 GPU 操作优先通过 CUDA stream 的提交顺序建立依赖,只有 GPU 结果需要交给 CPU 节点或服务返回响应时才进行必要同步。跨节点 Tensor 的数据搬运也由 Runtime 统一规划,避免每个算子独立管理跨设备传输。

2.3 固定 Beam Search 的单节点编译

固定 Beam Search 可以把有限次数的状态转移表达为结构确定的 Tensor 程序,同时保留由模型输出决定的候选内容。开放式生成系统通常需要在 decode step 边界重新组织请求状态并决定下一轮执行对象;当前负载的搜索深度固定,这部分控制因而可以从在线 Runtime 移入模型程序。

设最大生成 Token 数为 N_gen。模型首先执行 prompt prefill,并根据首步 logits 初始化候选;随后固定执行 N_gen - 1 次状态更新。每一步维护的主要状态包括:

  • 当前 logits;
  • 候选 token 和候选分数;
  • 累计 Beam 分数;
  • parent index;
  • 已生成 token sequence;
  • suffix K/V。

这些状态的内容由当前模型输出决定,但状态 Tensor 的容量、更新流程和算子组合在部署阶段已经确定。因此,Beam Search 不由 C++ Runtime 手写逐 Token 控制,而是在 Torch 模型中表示为固定结构的张量程序,并与 Transformer 一同导出和编译。

prompt prefill
  -> 首步候选初始化
  -> [
       decode
       -> 局部候选选择
       -> 累计分数合并
       -> parent index 更新
       -> token 与 suffix K/V 重排
     ] × (N_gen - 1)
  -> 最终候选及分数

整个过程最终封装为一个生成 AOTI 节点。DAG Runtime 对该节点只执行一次节点调度,不感知内部包含多少次 decode,也不理解 Beam Search 的候选选择和状态更新语义。

固定的是控制流,不是模型输出。parent index 仍会动态决定 suffix K/V、token sequence 和累计分数的重排,候选内容也仍取决于每一步 logits;这些数据依赖不会改变 Runtime 所看到的节点拓扑、Tensor 接口和资源规模。

将逐 step 控制收进编译节点后,请求不再在每个 token 后返回通用生成 Runtime。Transformer 计算、候选选择和状态更新可以在同一导出图中统一优化,在线服务只看到 generation 节点的输入和最终输出。

2.4 节点实现的离线编译

离线编译链的核心作用是将模型语义与硬件实现选择解耦。Python 侧负责表达模型和算子语义、选择 lowering 并构建产物;在线服务和 DAG Runtime 只加载已经生成的节点产物,不导入 Python 模型代码,也不进行即时编译。

对于具有独立数学语义、值得针对目标 GPU 选择实现的计算单元,框架定义稳定的可替换算子接口。Torch 参考实现明确规定输入输出、shape、dtype 和数学语义,构建配置再根据目标 shape 和硬件选择 Torch、Triton 或原生 C++/CUDA 实现。lowering 只替换节点内部的具体计算路径,不改变 DAG 可见的 Tensor 接口。

整体构建链如下:

Torch 模型
  -> 算子 lowering
  -> torch.export
  -> AOTInductor
  -> AOTI 共享库
  -> 节点配置与 manifest
  -> 可部署节点产物

Triton lowering 可以通过 torch.library.triton_op 和 wrap_triton 进入导出图;需要原生 C++/CUDA 扩展时,则在目标构建过程中将对应接口和静态库链接进节点产物。权重转换、packing、量化参数准备以及目标 GPU 相关的 kernel 配置也都在离线阶段完成。

模型定义还将运行时配置、编译配置和验收标准固化为可检查的部署描述:

  • 运行时配置:DAG 拓扑、节点 I/O、执行引擎、动态 shape 上界和并发资源;
  • 编译配置:模型参数映射、可替换算子的实现、目标 GPU 和精度策略;
  • 验证标准:参考实现、代表输入、边界 shape、数值门限、业务指标和端到端性能门禁。

模型结构和节点接口不变的新 checkpoint 可以复用大部分构建描述,但仍需重新检查参数映射、数值一致性和性能。接口复用并不意味着性能配置可以直接继承。

这种分层还使局部实现能够在接口不变的情况下替换。候选 kernel 先通过算子基准和数值对照筛选,再进入节点和完整 DAG 的 A/B 测试;只有在完整在线链路中仍能形成可辨认收益的实现,才会进入最终部署配置。由此避免仅根据 microbenchmark 决定线上实现。

明确的语义参考、可替换 lowering 和分层验收也为 AI 辅助优化提供了边界:AI 可以参与模型重表达和候选实现生成,最终选择仍由确定性的数值、业务指标和端到端性能验证决定。

2.5 执行资源与在线准入

状态和控制流有界,使执行资源能够在模型加载阶段按并发配置预先建立。在线链路中存在三个不同层次的并发参数,三者约束的资源不同,不能合并为一个统一的 concurrency 配置。

DAG 请求上下文持有一次请求执行所需的预分配 host/device buffer、逻辑 shape 状态和 CUDA stream。AOTI 执行实例负责调用已经加载的模型产物,并由执行引擎管理 runner、workspace 和权重资源。

加载阶段根据目标 GPU、动态 shape 上界和配置并发建立 DAG 请求上下文、Tensor 物理 buffer 和 AOTI runner。请求热路径只获取可用上下文、写入有效输入、更新逻辑 shape 并提交节点任务。

系统采用单个请求不跨卡、每张 GPU 部署完整模型副本的方式,不涉及 Tensor Parallel 或 Pipeline Parallel。不同卡型的算力、显存带宽、功耗限制和软件环境可能不同,因此节点内 kernel、精度配置和执行并发均按 GPU 独立选择;DAG、节点接口和资源管理机制保持不变。

容量配置按照完整链路逐层确定:先固定节点实现,再扫描 AOTI 执行实例数和 DAG 请求上下文数,最后设置服务入口的在途请求上限。每一步均通过端到端 QPS 和时延复核,避免将节点内部最优并发直接视为服务最优并发。

服务入口在已经验证的执行容量之外承担过载保护。后端暂时满载时,请求可以在入口有限等待;超过业务允许等待时间的请求不再进入后端计算。若持续到达速率高于单卡处理能力,服务按照配置拒绝或超时退出多余请求,避免已经失去业务时效的请求继续占用执行资源。

通过上述划分,在线系统只负责按照 DAG 执行完整请求,模型语义、固定生成过程和硬件实现均在节点构建阶段确定。第 3 章进一步介绍节点内部针对变长 prefill、Beam 状态、矩阵计算和 CPU 数据路径的专项优化。

03 有界生成节点的计算与数据路径优化

完成执行边界的调整后,系统吞吐主要取决于模型计算节点和相邻 CPU 节点的处理能力。两类负载中仍存在五类主要开销:变长样本的 padding 计算、Beam 间重复的 prompt K/V、候选选择与状态更新、密集矩阵计算,以及文本处理和结果编码的固定开销。

这些优化均以节点接口和数学语义不变为前提。Torch 实现提供数值和语义参考,构建阶段再根据目标 shape 和 GPU 选择具体 lowering 策略;候选实现依次通过算子、AOTI 节点和完整 DAG 验证后进入部署配置。

3.1 变长样本的 Packed Prefill

Prefill-Only 请求内包含多条长度不同的样本,统一 padding 会使所有样本按照最长序列执行线性层、MLP 和 attention。系统从 tokenizer 开始保留序列边界,将有效 Token 连续写入 packed buffer,并在一个模型节点中完成整个请求的 prefill。

请求内的多条字符串
  -> 连续 UTF-8 payload + lengths
  -> 批量 normalize / tokenize
  -> packed token buffer + sequence boundary
  -> packed Transformer prefill
  -> 各序列末位置 hidden state
  -> 每条样本的类别或分值

设第 i 条序列的长度为 L_i,query head 数为 N_q,K/V head 数为 N_kv,head dimension 为 D。packed attention 使用以下输入输出布局:

Q:   [ΣL_i, N_q,  D]
K/V: [ΣL_i, N_kv, D]
O:   [ΣL_i, N_q,  D]

sequence boundary 隐式定义 block-diagonal causal mask:不同序列之间不存在 attention 连接,每个 Query 只能访问同一序列内不晚于自身的位置。与统一 padding 相比,线性层和 MLP 的序列维计算规模由 B × L_max 降至更接近 ΣL_i,attention 则由 B × L_max² 降至更接近 ΣL_i²。batch 内长度差异越大,减少的无效计算越多。

FlashAttention-2 已提供成熟的变长 attention 实现。但在项目部署环境中,其变长接口无法同时满足 AOTI 集成与目标 GPU 兼容要求。系统因此采用针对目标 shape 优化的 Triton Packed GQA。

Triton kernel 不展开 GQA 的 K/V head,也不物化完整 attention score。kernel 根据 sequence boundary 确定当前 query block 对应的有效 K/V 范围,对 K/V 分块扫描,并使用 FP32 累积 online softmax 的最大值、归一化因子和输出状态。模型图随后使用同一套 sequence boundary 提取各序列末位置的 hidden state,再计算标签 logits、业务分值或生成所需状态。

sequence boundary 从 tokenizer 贯穿 attention 和结果提取,并与模型一起编译进 AOTI 节点。模型节点由此可以一次处理请求内的完整 batch,并按样本边界恢复输出。

Triton Packed GQA 使用显式 Torch 实现作为数值参考。代表 shape、边界 shape 和多序列用例共同验证数值误差与序列隔离。相对显式 Torch Packed Attention,该实现使中等和长 prompt 负载的完整 DAG 吞吐分别提升约 12% 和 39%。

3.2 请求内共享前缀的 K/V 布局

有界 Beam Search 中,所有 beam 共享同一个 prompt,仅生成 suffix 随搜索路径发生变化。这种固定的共享关系允许系统只保存一份 prompt K/V,并为每条 beam 保存较短的 suffix K/V:

prefix K/V: [L_p,    N_kv, D]
suffix K/V: [W, L_s, N_kv, D]

query:      [W,      N_q,  D]
output:     [W,      N_q,  D]

若每条 beam 分别保存完整 K/V,状态规模与 W(L_p + L_s) 成正比;使用 shared-prefix 布局后,则与 L_p + W L_s 成正比。对于本文固定 beam 宽度、同步推进且生成深度较小的负载,prompt 状态只需建立一次,请求结束时再与各 beam 的 suffix 一并释放。

对第 b 条 beam,attention 的数学语义仍为:

Ob=Attention⁡(Qb,[prefixK;suffixK[b]],[prefixV;suffixV[b]])O_b = \operatorname{Attention} \left( Q_b, [\mathrm{prefix}_K;\mathrm{suffix}_K[b]], [\mathrm{prefix}_V;\mathrm{suffix}_V[b]] \right)

其中分号表示序列维拼接。实际执行不会生成对应的 concat Tensor。Triton kernel 直接读取公共 prefix 和当前 beam 的 suffix,并在同一组 FP32 online softmax 状态中合并两段 K/V 的贡献。

kernel 按 (beam, K/V head) 划分 program,每个 program 处理该 K/V head 对应的多个 Query Head。该路径不需要展开 K/V head,也不需要 page table 查询或 prefix/suffix 中间拼接。parent index 更新时,只重排较短的 suffix K/V、历史 Token 和累计分数,公共 prompt K/V 始终保持不变。

Shared-prefix 布局保留每条 beam 对完整 prompt 的 attention 计算,只改变 K/V 的物理组织和访问路径。收益来自减少 prompt K/V 副本、拼接 Tensor、分页索引和相关数据搬运。

项目曾评估 FlashInfer 的共享 page table Paged Decode 路径,该候选在当时的对照中使完整 DAG 吞吐提升约 3%,但它与 FlashAttention 路径在集成至 AOTI 编译流程后均存在稳定性问题,因此未被采用。

3.3 Beam 候选计算与状态更新

固定 Beam Search 沿用原有的两级候选选择规则:每一步先处理 W × V 个词表分数,再从各 beam 的局部候选中选出全局候选。系统保持候选集合和累计分数语义不变,将候选选择、精度转换和状态重排改写为固定 Shape 的 Tensor 数据流,并且只搬运下一步真正保留的状态。

设词表大小为 V,beam 宽度为 W。完成首步候选初始化后,每个 decode step 会生成 [W, V] 的 logits。系统先在每条 beam 内选出 2W 个局部候选,再将这些候选与历史累计分数相加,最后从 W × 2W 个扩展中选出全局候选:

BF16 logits [W, V]
  -> 每条 beam 选出 2W 个 token ID
  -> 从 FP32 log-probability 中读取对应分数 [W, 2W]
  -> 加上各 beam 的历史累计分数
  -> 展平为 W × 2W 个候选
  -> 选出全局候选及其 parent index
  -> 仅为下一轮 active beam 重排 token、分数和 suffix K/V

中间步骤保留 2W 个全局候选,是为了同时接收已经生成 stop token 的完成序列,并从其余候选中继续选出 W 条 active beam。局部裁剪不会改变这组全局候选:任意一条 beam 排名在本 beam 第 2W 名之后的 Token,都不可能进入全局前 2W。因此,全局候选合并处理的数据量从 W × V 降为 W × 2W;当词表远大于 beam 宽度时,这一差异十分明显。裁剪边界存在同分候选时,具体 Token 和顺序由冻结的 Top-K 实现决定;业务语义只要求分数等价及相同实现重复执行结果一致,不依赖同分候选的固定顺序。

候选 ID 和候选分数采用不同的数据路径。log-softmax 对同一行 logits 保持严格单调,因此 Token ID 可以直接在 BF16 logits 上通过 Top-K 确定;候选分数则由完整词表上的 FP32 log-softmax 计算,再按照入选 ID 读取。这样既保留累计概率所需的 FP32 精度,也避免 Top-K 再扫描一份 FP32 词表 Tensor。

全局候选展平后,一个索引即可同时确定来源 beam 和该 beam 内的候选位置。来源 beam 被记录为 parent index,下一步直接用它选择对应的历史 Token、累计分数和 suffix K/V。已经淘汰的候选不再复制搜索状态,公共 prompt K/V 也不参与重排。

上述候选裁剪、分数合并、parent index 计算和状态重排均位于固定 Beam Search 的 AOTI 节点内部。每个 decode step 只执行定长 Tensor 操作,不返回在线 Runtime 创建候选对象或重新组织请求。项目也评估过专用 GPU Top-K 实现,但完整 DAG 没有获得稳定的额外收益,最终保留 Torch Top-K,并将优化集中在候选数据流和整体编译上。

3.4 密集矩阵计算的选择性 FP8

代表模型的大部分矩阵 FLOPs 集中在 MLP、其他线性层和 LM head。系统优先将这些计算密集且能够形成端到端收益的路径转换为 FP8,同时在 softmax、候选排序、累计分数和最终业务分值等敏感位置保留 FP32 中间计算。

以 MLP 为例,内部数据路径为:

BF16 activation
  -> per-token dynamic E4M3 quantization
  -> FP8 Gate/Up GEMM
  -> FP32 SiLU(Gate) × Up
  -> per-token dynamic E4M3 quantization
  -> FP8 Down GEMM
  -> BF16 output

activation 使用 per-token 动态 scale,权重在离线阶段转换为 E4M3,并保存 per-output-channel 的 FP32 scale。节点之间仍传递 BF16 Tensor,使低精度计算只局限在已经通过验证的矩阵路径内。

Gate/Up projection 通过一次 FP8 scaled matmul 产生中间结果。随后的 Triton kernel 以每个 token 对应的一行为处理单位,在一次遍历中完成 FP32 SiLU、逐元素乘法、amax reduction、scale 计算和 E4M3 重新量化,避免生成完整的 BF16 SwiGLU 中间 Tensor。Down projection 直接读取量化后的 activation 和 per-token scale,再通过第二次 FP8 scaled matmul 返回 BF16 输出。

动态量化、scale 计算和格式转换本身也会占用时间,较小或计算占比较低的 linear 无法稳定获得收益。长 prompt 模型最终只对 MLP 使用 FP8;短 prompt 模型的 LM head 计算占比更高,在通过数值和业务验收后也纳入 FP8。以 Prefill-Only 模型为例,MLP selective FP8 相对全 BF16 路径使完整 DAG 吞吐提升约 24%。

attention 点积累积和 online softmax 状态、log-softmax reduction、动态量化 scale、Beam 累积分数以及最终业务敏感分值继续保留 FP32 中间计算。这些精度保护策略并非完全通用。不同模型需要分别确定 FP8 覆盖范围,kernel 和分块配置也需要针对目标 GPU 固化;直接沿用其他模型或卡型的配置,可能同时损害业务效果和性能。

3.5 CPU 文本处理与结果编码

在小模型和短输出负载中,tokenizer、字符串处理和结果编码难以被 GPU 计算摊薄,优化前已经成为完整请求时延中不可忽略的组成部分。这些固定执行且数据规模有界的处理统一由 C++ CPU 节点完成。

业务 tokenizer 在基础词表之外增加了一批需要优先识别的 added token。原有实现会对这些 Token 逐项执行字符串查找;当前实现把 Token 的 UTF-8 字节组织成 byte trie,即按公共字节前缀建立的树形索引,使输入可以在一次扫描中完成匹配。匹配仍遵循“优先选择最靠左的起点,同一起点选择最长 token”的既有语义。

服务首先把一个请求中的多条字符串转换为一段连续 UTF-8 字节和每条字符串的长度,CPU 节点据此确定各条文本的边界。文本规范化和 tokenizer 随后直接写入请求上下文中预分配的 token buffer,并同时生成 packed attention 所需的 sequence boundary,避免为每条文本反复创建字符串对象和临时数组。

string input
  -> contiguous UTF-8 payload + lengths
  -> normalize / tokenizer
  -> packed token buffer + sequence boundary
  -> GPU model

结果侧同样使用预分配的连续 buffer。CPU 节点一次解码同一请求中的全部 Beam 候选,并把 Token、Score 和 length 直接写入预分配的 UTF-8 JSON buffer。buffer 的物理容量按最大输出预先建立,实际 length 标记本次请求写入的有效字节数;在线服务据此构造最终 RPC response。

Beam token / score Tensor
  -> 同一请求内批量 decode
  -> UTF-8 JSON payload + lengths
  -> RPC response

这条路径减少了请求热路径中的对象创建、动态内存分配和格式转换。byte trie 与预分配 tokenizer 将代表 prompt 集合上的文本处理由毫秒级压缩到十微秒级。

04 完整链路的验证与性能结果

系统验证沿算子、AOTI 节点、完整 DAG 和在线服务逐层展开。正确性测试确认各层实现保持相同的模型与业务语义,性能测试则以完整 RPC 的 QPS、平均时延和 P99 时延选择单卡容量点,再用分层 A/B 和有效 MFU 解释主要收益与剩余瓶颈。

4.1 测试环境与容量口径

性能评估在业务常用的相近算力档位单卡上完成,后续结果统一按这一档位的单卡容量呈现。每个模型均针对实际部署环境独立构建节点产物,并选择对应的 kernel、精度策略和执行并发。混合精度 MFU 按部署 GPU 的 dense 理论峰值计算,BF16 和 FP8 峰值分别取约 150 TFLOPS 和 300 TFLOPS。

容量测试采用固定在途请求数的 closed-loop 方式。每个并发档位始终保持相同数量的完整请求正在等待或执行,一个请求完成后立即发送下一个,中间不加入客户端等待时间。QPS 统计每秒完成的完整 RPC 数量,端到端时延从请求发出持续到结果返回,P99 表示 99% 的请求不超过该时延。逐档增加并发可以得到单卡持续处理能力和时延开始快速增长的拐点。

单卡容量取同时满足平均端到端时延低于 20 ms、P99 低于 80 ms 的最高完成 QPS。服务入口再通过 deadline、超时退出和拒绝策略阻止过量请求继续占用后端资源。

4.2 分层正确性验证

执行链路一致性。在相同机器、相同输入和相同部署产物下,Python Artifact Runtime 与 C++ DAG Runtime 的输出 Tensor 保持逐字节一致。Server 使用同一 DAG Runtime 执行请求,协议层继续核对输出名称、类型、有效长度和最终业务结果。这一层验证用于确认节点调度、buffer 管理、数据搬运和协议转换没有改变部署产物的语义。

专用计算的数值验证。Packed GQA、shared-prefix attention、FP8 linear 和其他自定义 lowering 分别与 Torch 参考实现进行算子级和模型级数值对照。Beam Search 进一步逐 step 比较候选 Token、候选分数、parent index、累计分数、suffix K/V 重排和最终序列。

业务指标验收。低精度版本与原始离线模型使用相同的数据集、样本处理、评估代码和指标定义。优化实现不要求与原始 BF16 模型逐元素一致,但业务指标变化需要处于预先确定的验收范围内。最终部署产物通过了上述各层验证。

备注:【分层正确性验证】专用计算先与 Torch 参考实现进行数值对照;同一部署产物从 Python Artifact Runtime 到 C++ DAG Runtime 保持逐字节一致;低精度版本最终使用相同的离线评估流程完成业务指标验收。

4.3 时延约束下的单卡容量

下表给出三类代表负载在完整 RPC 链路上的单卡容量。三类模型的参数规模约为 0.5B 至 0.6B。每个 Prefill-Only 请求包含十余条样本,因此同时报告请求吞吐和样本吞吐;Beam Search 按一次完整搜索统计请求吞吐。

容量扫描表明,无约束峰值吞吐并不等于满足业务时延目标的容量。以短 prompt Beam Search 为例,在选定容量点上继续增加执行并发后,QPS 仅增加几个百分点,平均端到端时延却上升 30% 以上,并超过 20 ms 目标。

从并发曲线看,较低并发下增加请求上下文可以提高 CPU 处理、kernel launch、数据传输和 GPU 计算之间的执行重叠;进入饱和区间后,继续增加并发主要转化为排队、资源争用和单请求执行时间增长。线上配置因此选择满足时延约束的容量点,而不是无约束条件下的最高 QPS。

4.4 关键优化的分层 A/B

不同优化需要在对应层级进行 A/B,并最终回到完整请求确认收益。下表汇总 CPU 节点、AOTI 节点和完整 DAG 层级的主要结果。

分层 A/B 用于判断局部优化能否转化为完整请求收益,并确认优化后瓶颈所在的位置:

  • Packed GQA 相对显式 Torch Packed Attention 的收益随 prompt 变长而扩大,说明变长 attention 实现是长文本路径的重要瓶颈;
  • shared-prefix attention 通过减少 K/V 副本、拼接 Tensor、分页索引和数据搬运,使长 prompt Beam Search 的完整 DAG 吞吐提升约 20%;
  • 选择性 FP8 在保留敏感计算精度的同时,使 Prefill-Only 的完整 DAG 吞吐提升约 24%;
  • tokenizer 的文本处理由毫秒级压缩到十微秒级,CPU 热路径不再占用主要时延;
  • 在固定 1,024-token 且不包含 RPC 的诊断负载下,完整 DAG 与 AOTI-only 的吞吐差异约为 1%,说明这一测试范围内的框架附加开销较小。

随着模型主体计算被压缩,profile 显示瓶颈逐步转移到 attention 状态布局、低精度转换、CPU 数据路径、同步和结果编码等环节,因此每次局部优化都需要回到完整请求时间线中复验。

4.5 混合精度有效 MFU

为估算逻辑矩阵计算相对于硬件理论能力的占用程度,本文定义“混合精度有效 MFU”。该指标将每个请求中不同精度的逻辑矩阵 FLOPs 分别除以对应的 dense 理论峰值,再按照相同请求口径下的完成 QPS 汇总。

设每个请求中的 BF16 和 FP8 逻辑矩阵计算量分别为 F_BF16 和 F_FP8,对应硬件理论峰值分别为 P_BF16 和 P_FP8,则:

Effective MFU⁡=QPS(FBF16PBF16+FFP8PFP8)\operatorname{Effective\ MFU} = QPS \left( \frac{F_{\text{BF16}}}{P_{\text{BF16}}} + \frac{F_{\text{FP8}}}{P_{\text{FP8}}} \right)

公式先估算每个请求中 BF16 与 FP8 逻辑矩阵计算在理论峰值下所需的等效时间,再乘以相同请求口径下的完成 QPS。Prefill-Only 的一个请求包含完整 batch,Beam Search 的一个请求包含一次完整搜索。结果表示单位 wall-clock 时间内完成的逻辑矩阵工作量相当于理论峰值算力的比例。

MFU 统计 linear、LM head 和 attention 的逻辑矩阵 FLOPs,FMA 按 2 FLOPs 计数。以下工作不进入分子:

  • 物理 shape 对齐或分块产生的冗余计算;
  • 动态量化和 scale 计算;
  • 逐元素计算与非线性;
  • softmax 和 log-softmax;
  • Top-K 与 Beam 状态更新;
  • tokenizer、结果编码和 DAG 调度;
  • 数据传输和同步。

代表结果如下:

计算使用对应部署 GPU 的 BF16/FP8 dense 理论峰值,取约 150/300 TFLOPS。

该指标用于判断后续优化方向。Prefill-Only 模型的有效 MFU 已较高,继续只优化 GEMM 的端到端边际收益将受到限制;短 prompt Beam Search 则需要继续关注 Top-K、搜索状态和短 shape 计算。

05 适用范围与结论

5.1 框架与节点的复用边界

只要完整请求处理流程能够拆分为一组接口明确的节点,并由 DAG 表示节点之间的数据依赖,就可以接入本文所实现的这套推理框架。框架统一提供节点调度、Tensor 搬运、设备同步、资源预分配、请求上下文和并发管理,使新模型的开发与性能分析首先聚焦节点内部计算。

框架能力和专项优化具有不同的复用范围:

  • 框架级能力:DAG Runtime、执行引擎、资源生命周期、并发管理和链路一致性检查,可由不同模型直接复用;
  • 节点级语义:Packed Prefill、固定 Beam Search 和 shared-prefix attention,可由数据布局与状态结构相同的模型复用;
  • 节点级实现:具体 Triton/CUDA kernel、FP8 覆盖范围、分块配置和执行并发,需要按照目标 shape、GPU 和软件环境重新选择。

模型结构和节点接口不变的新 checkpoint 还可以复用大部分构建描述,再完成参数映射、数值、业务指标和端到端性能验证。静态 DAG 本身不限于有界生成负载,本文的专项节点与性能结论则面向 Prefill-Only 和固定深度 Beam Search。

5.2 专项优化的适用边界

本文的专项优化主要面向控制流和状态容量稳定、单次生成较短的 SLM 负载。对于以下场景,通用 LLM serving 系统所提供的动态能力与需求更直接匹配:

  • 开放式长文本生成;
  • decode 长度高度动态;
  • sequence 持续进入和退出;
  • 强依赖 continuous batching;
  • 跨请求公共前缀较长且命中率较高;
  • KV 状态规模较大,需要动态分页、淘汰或复杂内存复用。

在这些场景中,vLLM、SGLang 等系统提供的 continuous batching、paged KV cache、prefix cache 和动态请求管理更匹配负载需求。当前部署采用单个请求不跨卡、每张 GPU 部署完整模型副本的方式。

5.3 结论

搜索推荐中的部分 SLM 保留 Decoder-only Transformer 和 next-token prediction 结构,但在线执行只包含一次 prefill 或固定次数 decode。请求规模、Beam 宽度、生成深度和状态容量的上界均能在部署前确定,因此模型构建、资源规划和在线执行可以围绕完整业务请求重新组织。

系统使用静态 DAG 表示完整请求流程,以 CPU 自定义节点处理文本和结果,以 AOTI 节点封装模型计算与固定搜索过程。DAG Runtime 负责节点调度、Tensor 搬运和请求资源,节点内部则分别应用 Packed Prefill、shared-prefix K/V、选择性 FP8 以及候选与状态更新优化。这种分层将通用在线执行能力与模型专项优化解耦,也使正确性和性能能够逐层验证。

在本文覆盖的模型、输入分布和单卡环境中,Prefill-Only 达到约 100 request QPS,即约 1,500 sample/s;短 prompt 和长 prompt Beam Search 分别达到约 300 和 120 request QPS。三类负载的平均端到端时延均低于 20 ms,P99 约为 70 ms,满足 80 ms 的线上要求。相同部署产物从 Python Artifact Runtime 到 C++ DAG Runtime 保持逐字节一致,低精度路径与专用算子的业务指标变化也处于预设验收范围内。Prefill-Only 与短 prompt Beam Search 代表负载的混合精度有效 MFU 分别约为 67% 和 41%。

这组结果表明,推理系统的设计应由实际执行边界决定。对于有界生成负载,将固定结构前移到构建阶段、按完整请求组织在线执行,并把专项优化限制在接口明确的节点内部,可以提高计算效率和时延约束下的单卡容量。



Source link

Leave a Reply

Your email address will not be published. Required fields are marked *