RTP-LLM 在相关性大模型中的推理优化最佳实践

在淘宝搜索场景下，用户Query与候选商品(Item)之间的相关性判别是非常重要的一环，它筛选出该Query下最相关的商品， 是用户体验的基石。过去几年主搜在相关性场景上已经做了不少工作，并且取得了显著的正向收益，今年，为了进一步解决部分口语化Query承接效果较差的问题，我们又引入了更大参数量(激活3.5B的MoE)的LLM模型，同时扩展了单次打分商品个数，这给我们的系统性能提出了巨大挑战。主要体现在：

• 超大规模Point-wise计算：针对每一个Query请求，相关性模型需要对多个(Query, Item)对进行相关性判断，计算量随候选集大小线性增长。

• 超长Context的Prompt处理：为了让大模型准确理解商品细节，我们需要将 Query 和包含丰富商品属性、描述信息的 Item 拼接成一个很长的 Prompt 输入给模型，长度带来的是O(N2)计算量增长。

• 极度严苛的RT要求：搜索端到端的延迟要控制在1秒内，抛出Query理解、召回和客户端传输解析环节，留给相关性模型的空间也就在500ms左右。

在大BatchSize和超长输入的双重压力下，如何保证大模型推理满足如此严苛的时延标准，是最大的难题。本文将重点分享我们如何基于RTP-LLM框架，通过一系列技术创新，成功支持LLM相关性模型成功上线。

在LLM相关性计算场景中，面对单次请求需处理 多个 (Query, Item) 对的超大输入，如何高效调度大模型推理资源，在有限的时延约束下完成计算，是工程落地的首要难题。本章将重点介绍我们在 Proxy 调度层面的关键技术演进。

相关性模型选用的是激活3.5B的MoE模型，直接发送全部请求给单机推理，响应时间会长达数秒，这对于主搜是绝对不可接受的。为了降低时延，我们进行了横向扩展，通过拆批并行计算加速。假设下游有N个推理节点，单次推理M个(Query, Item)对，理论上，将M个pair拆为n个小批次(n<N)，每批包含(M/n)个pair并行发送，理想情况下RT可降低至原来的1/n。然而，传统的服务发现与负载均衡组件(如VIPServer)通常基于请求粒度进行随机分发，无法保证瞬间并发子请求的均匀，导致下游部分机器出现排队，拉高了整体RT；

为了解决这一问题，我们重新设计了 Proxy 的调度策略，实现了精细化的拆批与动态负载均衡：

1. 基于计算量的均匀拆分： 我们摒弃了简单的固定Batch Size拆分方式，而是基于Query /Item Token预估长度动态组合，确保每个批次的预估计算量基本一致，避免因任务倾斜导致的某个 Batch 计算过慢。

2. 自定义动态调度策略： 我们引入了专用的 Proxy 服务，该服务会实时维护每个下游机器的活跃请求数和上一次请求的完成状态。当新的大请求到达时，Proxy 会基于实时负载信息，采用加权轮询或最小连接数等策略，优先将拆分后的小 Batch 发送给当前最空闲的可服务机器，从而实现整体 RT达到全局最优。

确定了拆批策略后，如何配置单个推理实例的并行策略以获得最佳性能，是另一个关键决策点。以BS=16为例，我们对比了2种部署方式：

• 方案 A (4 * 1TP)： 将 BS=16 进一步拆分为 4 组BS=4，发给4台独立的单卡机器。

• 方案 B (1 * 4TP)： 将 BS=16 作为一个整体，发送给一台 4 卡张量并行机器。

测试表明，如下图timeline所示，方案B受限于模型结构(MoE导致的专家节点负载不均衡)以及TP模式下通信开销，当出现个别慢节点时，整个Batch的完成时间会被拖慢；相比之下，方案A中各实例相互独立，无通信开销，可以更好的利用数据并行性。在实际测试中，BS=16的情况下，方案A比方案B降低12ms，成为线上实际部署模式。

起初，为了快速迭代和验证想法，Proxy 服务采用了 Python 实现。然而，在实际部署上线后，受限于Python在高并发场景下固有的性能瓶颈和GIL 及 GC 机制，经常出现剧烈毛刺，严重影响了服务的稳定性。为了解决这一问题，我们将Proxy 服务迁移至天生支持高并发的 Go 语言重构。迁移完成后，效果立竿见影：服务毛刺大大缓解，在晚高峰时段，P99 延迟从原先的 800ms 大幅下降至 500ms，降幅达 300ms，有力地保障了在线服务的稳定性。

在相关性场景中，模型的输入 Prompt 具有显著的结构化特征：同一个 Query 会与多个不同的 Item 进行组合评估。在一个 Batch 内的所有请求中，Prompt 的 Query 部分的 Token 是完全相同的。在标准推理流程下，这部分公共前缀的 Attention 计算会被重复执行，造成了算力浪费。

RTP-LLM中的原有的 KV-Cache 复用技术是基于批次间的。只有当前一个推理请求完全结束，其对应的 KV Block 才会写入 Cache Manager 供后续请求复用。这种机制无法解决我们面临的单批次内请求的公共前缀冗余计算问题。

一种朴素的思路是将 Query 前缀提取出来进行预推理，然后再发送包含完整 Item 的请求。然而，这种“两阶段推理”方案存在显著缺陷：一方面增加了 Proxy 的复杂度，需要维护复杂的请求状态和一致性 Hash 路由；另一方面，额外的网络传输和推理引擎调度开销往往会抵消掉预计算带来的算力收益，导致实际 RT 不降反升。

为了从根本上解决这一问题，最好的方式就是实现推理批次内部的 KV-Cache 复用。为此，我们对RTP-LLM 的 KV-Cache 管理机制进行了改造。核心思路是引入一种投机式的 Block 分配与注册机制：

1. 调度阶段的预注册： 在 Scheduler 对批次内请求进行调度时，我们顺次为请求分配 Block ID。对于第一个请求，在分配完 Block 后，尽管此时Block内尚未填充真正的KV值，我们立即将其公共前缀 Token 与 Block 的映射关系注册到 Cache Manager 中。

2. 批次内复用： 当调度该批次内的后续请求时，调度器会查询 Cache Manager，发现公共前缀已存在映射关系，从而直接复用已分配的 Block ID，无需重新分配。

3. Lazy Filling： 在真正执行推理的 Attention 计算阶段，引擎会先行计算当前Token的KV值并将其更新填充到对应的Block中；随后，再从这些刚刚完成更新的Block中读取KV数据进行最终的Attention计算。GPU 单 Stream 内 Kernel 严格的顺序执行保证，可以保证结果的正确性。

上述机制引入了一个严峻的 Cache 一致性挑战：在 KV 数据尚未真正计算完成时就提前注册映射关系，一旦该次推理中途失败或中断，Cache Manager 中就会残留指向未初始化或错误数据的“脏 Block”，进而污染后续所有复用该 Block 的请求结果。

为了解决这一问题，我们引入了一个内部标识Epoch来管理 Cache 的可见性范围。

• 隔离机制： 当一个批次正在推理时，其预注册的 Block 被标记当前批次的Epoch，该 Block 仅对同一批次内的其他请求可见，对全局其他批次的请求不可见。

• 事务提交与回滚： 只有当该批次推理成功完成，这些 Block 的状态才会被更新为全局可见，完成“提交”。反之，如果请求调度失败，系统会触发回滚操作，自动清理未提交的 Cache 映射并释放 Block 资源。如果推理中途失败，该Cache映射会一直处于不被命中状态会被CacheManager释出。从而杜绝了 Cache 污染的风险，在确保极致性能的同时保障了系统的稳定性和结果正确性。

以上批次内前缀复用，在相关性场景获得了10%的收益，理论上可以通过重训算法模型扩大复用的范围，不仅仅是Query侧的输入，还有Item侧相同的属性等特征提前，还可以进一步扩大收益。

后续我们注意到，在VLLM 1.0中，同样支持了批次内KV-Cache复用。这种技术路线上殊途同归，说明消除前缀冗余计算是业内的共识。

相关性模型采用了MoE架构。通过性能分析我们发现，MoE计算占据了前向推理总耗时的70%以上，成为了主要的性能瓶颈。提到MoE的优化，往往会想到多卡多节点的EP方案，但主搜这种在线服务场景，它追求的不是极致吞吐量，反而对推理latency有严苛要求。在这个限制下，我们选择的单卡小批量推理方案虽然满足了延迟的要求，但也带来了专家激活的稀疏性问题，给计算效率带来了严峻挑战。

在小批量推理模式下，输入到MoE层中的Token总数有限，经过路由分配后，每个专家实际处理的Token个数变得稀疏。我们对线上真实场景进行了采样分析(Batch Size=4，平均序列长度576)，下图是各专家的Token分配情况：

可以看出，绝大多数专家分配到的Token数量都小于64，只有极个别专家的Token数超过128。这种稀疏的负载和底层执行机制产生了冲突，因为我们采用的MoE后端是DeepGemm，其默认的计算粒度blockM是128。实际计算时，为了充分利用Tensor Core进行高效的矩阵运算，DeepGemm会将每一个专家分配的Token数padding到128的整数倍。这种方式在当前场景引入了大量的无效计算，造成了算力浪费。

针对上述问题，我们引入了动态kernel选型机制，根据实时的输入负载动态选择最优的Kernel，从而最小化Padding带来的开销。该策略的核心逻辑如下：

1. 默认选择高效Kernel：Dense情况下，blockM=128的Kernel是效率最高的选择。因此，会优先选用block128 Kernel。

2. 引入Padding开销率：在执行计算前，系统会预先计算在blockM=128的策略下Padding 后的总 Token 数与原始总 Token 数的比值，我们称之为“Padding 开销率”。

3. 动态决策与切换：我们设定一个超参数阈值时(实验确定为1.5)，当Padding开销率大于该阈值时，则认为当前输入负载过于稀疏，会自动切换到粒度更小的blockM=64的Kernel进行计算，减少无效计算量。

通过这个动态调优策略，我们提升了MoE层的计算性能，在端到端上获得了可观的延迟收益。从以下实验结果可以看出，该优化在小批次下性能提升更明显，这也符合我们的预期：批次越小，Token分配越容易出现稀疏现象，动态调优优势也更明显。

尽管出于RT的限制我们没有选择EP，但如果我们能解决通信问题，EP仍然是一个诱人的方案。在传统的LLM推理中，因为专家权重单张卡放不下，EP往往会将多个expert的权重分布在不同的卡上，并且通过EPLB来平衡各张卡之间的计算量。而我们的场景下模型尺寸相对较小，可以在每张卡上都驻留全部专家，实时地根据token分布进行卡间的负载均衡，使每张卡上的计算量相当，同时需要发送的token最少。

批次内前缀复用还有一个可以改进的优化点：以kv cache page为单位的复用会导致不足一个page的前缀不在复用范围内。例如公共前缀的长度为90，page size为64，那么会有64长度的前缀被复用，而剩下的90-64=26个token则仍然需要被反复计算。我们可以写一个专用的attention来避免重计算这部分，充分利用相同的前缀。我们尝试做了一个demo，结果显示，还有额外的端到端3%-5%的优化潜力。

我们重点通过Proxy负载均衡/批次内KV-Cache复用/MoE Kernel动态调优等手段，成功支撑了40A3B LLM模型在搜索在线链路的落地，取得了相关性体验的全面提升。但这仅仅是个开始，面对日新月异的模型技术演进和不断提升的业务期望，我们将持续在性能、效率和成本上持续进行架构演进，持续探索搜推场景推理型模型的能力上界。