大模型推理框架 RTP-LLM Embedding 技术揭秘

2025-03-20
浙江
本文字数：4274 字
阅读完需：约 14 分钟

一、背景介绍

Embedding（嵌入）是现代机器学习和深度学习的重要组成部分，通过将离散数据映射到连续向量空间，解决了高维稀疏性和语义表达的问题。它在自然语言处理、推荐系统、计算机视觉等领域有着广泛的应用。RTP-LLM 是阿里巴巴智能引擎团队自研的大模型推理加速引擎，作为一个高性能的大模型推理解决方案，它已被广泛应用于阿里内部，本文将介绍项目在 Embedding 框架上的实践和思考。

在我们的生产环境中，主要存在两种使用 Transformer 模型实时生成 Embedding 的场景：一类是部署在云服务器或者内部大模型服务平台的 Pytorch Huggingface 模型，用于计算 Embedding 或者进行重排/分类；另一类是搜推广场景，使用 Tensorflow 的 Bert 模型计算商品和用户的相似度。这两类场景性能表现都一般，因此我们希望能够提供一个解决方案，能够在部署方便的前提下，优化上述两种场景 Transformer Embedding 计算的耗时和吞吐，减少资源消耗。

二、问题描述

前文描述的两类场景性能不好的问题可以分成两方面，模型算子实现和外围请求调度策略。

在模型性能层面，Pytorch 和 Tensorflow 模型默认都没有使用高性能的 FlashAttention 算子，没有做算子融合，在模型参数大的情况下也没有量化支持。并且原生 pytorch 模型 launch 开销大，GPU 执行效率低[1]。

在外围请求层面，搜推广场景的请求基本都是大批次请求，在进入模型时在序列长度维度对齐，这种凑批方式会导致无用计算增加，进而增加时间；而在线上服务平台层面，单个请求批次太小，无法充分发挥 GPU 算力。

因此我们希望解决方案在提供高性能算子（FlashAttention，Elementwise 算子合并等），模型量化的同时，在调度层对大批次请求去除空洞、对多个小批次请求进行合并的方式，解决上述的问题。

三、方案调研

在设计方案时，vLLM 和 TensorRt-LLM 都没有支持 Bert 模型的 Embedding 功能，因此最初我们考虑的解决方案是 TritonServer+TensorRT Backend。TritonServer 使用 C++编写，同时提供 HTTP 和 GRPC 服务，支持不同请求间凑批和多种请求调度策略；TensorRT Backend 提供了 GPU 上的高性能算子，并且支持通过编译的方式从 Pytorch 代码转换 onnx graph 最后构建 TensorRT Engine。

这套方案在性能和自动化部署两个角度都非常符合我们的需求，但是存在以下几个方面的问题：

"差一口气"的 token 级别凑批

长度较短的请求在 GPU 上运行的瓶颈是显存带宽而不是算力，因此合并请求有助于提高算力资源的利用率，从而提高吞吐。而对于不同请求间的凑批，传统的做法是使用 padding。假设输入 tensor 的 shape 是二维的[batch, seq]，则在 seq 维度 padding 到最大值，在 batch 维度 concat。但是这样的做法增加了无用的计算量，对于 seq 维度差异大的场景显然不合算。

这个问题有解决办法，Triton Server 侧实现了 Ragged Batching[2]，将输入打平成一维做 concat 以避免 padding，这种凑批方式我们称之为 token 级别凑批。并且 TensorRT 的 Bert 模型支持使用 Variable Sequence Length[3]，通过这种方式构建的引擎支持 token 维度凑批的输入。但是这种引擎无法通过编译的方式构建，必须编写 TensorRT 插件代码，在产品上非常不灵活。

模型支持欠缺

在 Pytorch->ONNX->TensorRT Engine 的构建模式下，除 Bert 外很多模型支持的都不是很好。以 Layernorm 操作为例，由于 Layernorm 计算公式中有平方和的累加，因此在半精度下经常会越界导致结果 nan。

Tensorrt 的解决方式是识别 ONNX 图中的 LayerNormalization 算子，将这部分运算使用 fp32 精度来解决。但是这两年流行的大语言模型（Large Lanuage Model, 简称 LLM）很多使用了 Rmsnorm 代替 Layernorm，ONNX 里没有 Rmsnorm 的算子，所以 TensorRT 无法识别到这个结构，会在 FP16/BF16 精度下进行这部分网络的计算，导致模型输出错误。

跨设备支持不友好

TensorRT 仅提供了 Nvidia GPU 上的模型加速，如果希望在 AMD 或者 CPU 上进行高性能计算，则需要额外开发 Triton Server 的 backend，开发成本大。

基于上述的考虑，最终我们没有使用上述方案，而是在我们自研的大模型推理引擎 RTP-LLM[4]中开发了 Embedding 模块。RTP-LLM 集成了 Cuda/Rocm/Arm 的高性能算子，对主流大模型都有比较完善的支持，并且框架内天然支持 token 级别凑批，在性能和迁移性方面都解决了上述提到的痛点。在 RTP-LLM 中支持 Embedding 功能，我们可以更多的把开发重心放在如何便捷的部署各种类型的 Embedding 模型。

四、框架设计

在介绍 Embedding 框架之前，先简单介绍下 RTP-LLM 在大模型场景下的生成链路：用户请求在编码后会先进入调度器（Scheduler），调度器管理请求生成所需要的资源，并在每轮生成时选择一批请求到执行器（Executor）。执行器根据请求的生成阶段（Prefill/Decode）将请求组装成模型接受的格式送到模型层执行，最后将结果更新后送回调度器进行下一轮调度。相比 LLM 框架，Embedding 框架做了以下三处改动：

调度层：Emebedding 可以认为是大模型生成中的 Prefill 阶段，但是不需要保留 KVCache。针对这个特点，我们简化了 Scheduler 和 Executor 里的流程，并且允许请求不申请 KVCache，减少显存的占用。同时在 Embedding Scheduler 会在限制范围内尽量的进行请求凑批，在单条请求延迟在接收范围的同时提高服务吞吐。

后处理：LLM 的后处理流程通常是采样，但是 Embedding 的后处理流程多种多样。对于同一个模型，在不同的下游任务下，会使用不同的网络处理 Transformer 层输出。这些网络结构大部分定义在 Python Embedding 库中（Sentence Transformer / Flag Embedding 等），小部分作为里的 Python 文件附带在 checkpoint 文件中。由于后处理流程相比 Transformer 网络耗时占比非常小，因此 Embedding 框架中这部分网络仍在 Python 中构建，在运行时通过 pybind11 调用，并且支持用户使用 SentenceTransformer 配置自定义这部分逻辑。小部分性能敏感的场景则支持通过 c++重写后处理代码以达到更好的性能。

请求层：和生成式任务相比，部分 Embedding 任务的结果较大（[batch, hidden_size] / [batch, sequence_length, hidden_size]），使用 json 存储会产生较大的响应包，且使用 HTTP 协议在 json 序列化和反序列化都会有较大也会产生额外的开销，最终造成性能上的损失，因此对部分场景我们也支持了 RPC 协议访问，包括开源的 Grpc 和内部 Arpc 两种方式请求 Embedding，减少这部分开销。

最后我们实现的 Embedding 框架如下图：

五、模型量化

RTP-LLM 中对模型的量化有 WeightOnly INT8/INT4 量化、Dynamic SmoothQuant[5]和 FP8 量化，这些量化方法通过提升 Transformer 模型的矩阵乘速度的方式加速模型推理。在大部分 LLM 场景这些量化都有不错的表现，但是在搜推广场景上这些优化都没有收益，原因如下：

WeightOnly 量化仅适用于矩阵乘输入数据(activation)远小于权重(weight)的场景，此时矩阵乘是访存瓶颈任务，WeightOnly 量化通过降低模型权重大小的方式减少访存需求，从而加速矩阵乘。而搜推广场景使用隐乘大小是 768 的 Bert 模型，这时矩阵乘权重 shape 为[768, 768], [768, 2304]和[768, 3072]，而输入数据的大小平均是[3840, 768]，此时矩阵乘是计算瓶颈任务，而 WeightOnly 量化需要在运行时反量化增加了计算量，甚至会拖慢矩阵乘的速度。

Dynamic SmoothQuant 量化通过同时量化矩阵乘输入数据和权重成 INT8 数据格式，并使用 INT8 Tensor Core 的方式，同时加速了模型权重读取和计算的速度，对于矩阵乘本身有明显的加速。但是不同于在加载模型时就量化好的权重，将输入量化成 INT8 的过程是实时的，因此需要引入额外开销。实际测试我们发现在[3840, 768]的输入下，量化输入的新增时间基本和矩阵乘加速节省一致，因此 SmoothQuant 在这个场景下也无法发挥作用。

FP8 可以认为是 w8a8 量化的升级版，在保持较高精度的同时使用 FP8 Tensorcore 进行计算，而且框架还支持了 FP8 的 FlashAttention，但是搜推广场景线上大部分卡是 A10/V100，在硬件上缺少 FP8 的支持。

最后我们开发了 Static SmoothQuant 的方式加速 Bert 模型，Static SmoothQuant（下称 StaticQuant）和 Dynamic SmoothQuant（下称 DynamicQuant）的区别在对 DynamicQuant 操作中有一个 Reduce 操作统计数据的最大值，这一步会引入比较大的开销。

StaticQuant 将这一步放在训练时，通过统计训练数据中最大值近似，因此需要耗费的时间大大减少。并且 StaticQuant 的量化是一个 Elementwise 操作，可以通过算子融合的方式合并到 Gemm 和 Layernorm 算子中，进一步降低耗费的时间，具体前后对比如下：

我们以输入 shape 为[3840, 768]的情况下 A10 机器上测试，StaticQuant 模型单个请求耗时 2.9ms，对比原始模型 4.9ms 降低 40%。上线后量化模型对比原始模型效果损失在 0.1%以内，达到了预期效果。

六、总结与展望

我们基于 RTP-LLM 实现了 Embedding 框架，支持部署 Transformer 结构的 Embedding 模型及其下游任务（Reranker/Classifier），在请求上支持 HTTP/ARPC/GRPC 协议，在部署上支持用户使用 SentenceTransformer 自定义后处理逻辑。Embedding 引擎已服务了淘宝主搜等多个在离线场景，并成功度过双十一洪峰。虽然当前已解决许多问题，但未来仍有很多优化空间：

HTTP 服务的入口目前还在 Python 侧，这导致 HTTP 请求会受限于 Python 的性能，在高并发短请求时无法充分发挥引擎性能。在近期我们会改用 C++ HTTP 服务解决这部分问题。
由于后处理部分输入格式是完全确定的，可以引入 torch 的 compile 技术减少 kernel launch 消耗。
参考 vLLM 的 Multi-Step Scheduling[6]，将 Scheduler 和 Executor 改成是生产者消费者模式，并行化请求调度，数据传输和模型执行，进一步提高框架性能。

除此之外，我们也将尝试框架层之上的优化。比如目前我们使用的负载均衡策略是 RoundRobin，但是线上 Bert 集群的平均延迟和 p99 延迟仍有很大差距，原因是在机器上有并发的大批次请求，导致单个请求排队，是否有更好的策略能够在 Embedding 这种高 qps 短 rt 的场景让流量更加均匀呢？诸如此类的问题都需要我们继续探索和优化。

RTP-LLM 项目已在 GitHub 上开源，欢迎大家与我们交流共建！