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

一、背景介绍

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 维度差异大的场景显然不合算。

这个问题有解决办法，Trit