GR-Inference：让搜索生成式召回更高效
REDtech约 11 分钟基于 Semantic ID(SID) 的生成式召回，正在成为搜索侧解决长尾供给与结果多样性的重要建模范式。但这类负载与通用 LLM 服务差异很大：上下文很长、生成步长很短、Beam 很大。以本文的业务测试负载为例，每个请求只生成 3 个 token，但要同时维护并最终返回 900 条 Beam 候选。通用 Serving 框架在此场景并不能高效支持。本文介绍小红书与 NVIDIA 联合构建的 GR-Inference——一个专为“长 Context + 短 Decode + 大 Beam”设计的推理引擎，覆盖 KV 抽象、连续批处理、专用 Decode Attention、受限生成与 CUDA Graph 等全链路优化。在固定 Beam 900 的 SLA 口径下，GR-Inference 相对三套 SGLang Beam Search 社区实现取得 1.5x~2.5x 的峰值吞吐优势；在 Item-constrained 场景下相对 SGLang 提升 3.2x 以上。
搜索场景下，传统召回以 ANN 双塔匹配为主：将 Query 与笔记分别编码成向量后，用向量相似度召回候选。该范式存在先天局限，Query 侧和笔记侧是各自独立编码的，两者的交互信息没有充分参与匹配，召回精度也受限于向量相似度的表达上限；而 LLM 强大的语义理解与生成能力，在这方面还有很大发挥空间。
为此，小红书在搜索侧落地了一路 LLM 生成式SID召回，用“生成”取代“匹配”：离线阶段对笔记的文本表征做分层残差聚类，为每篇笔记生成多级 SID；在线推理阶段用 LLM 根据用户 Query 及本次查询的上下文、意图等信息，逐级生成目标笔记的 SID，再把多级 SID 映射回对应的笔记候选，从而完成召回。
Figure 1. 小红书生成式召回链路
LLM生成式SID召回的核心是自回归的 GR（Generative Retrieval）模型：在推理阶段使用LLM 把 SID 逐级生成出来，前级 ID 决定大致方向，后级 ID 在前级基础上进一步细化（即每一级生成的逻辑都依赖用户 Query 信息与已确定的前级 ID）。同时，为了兼顾最终召回候选笔记的数量与多样性，推理时需同时保留多条候选路径（即多个 Beam）。这就是本文的主角——GR-Inference 所服务的核心形态。
GR 场景的推理负载有如下特点：
长 Context：一次请求要携带很长的用户历史与候选上下文（1k~5k Token 常见），Prefill 计算重；
短 Decode：真正生成只有 3~5 步（每级 SID 一步），Decode 阶段很短；
大 Beam：为保留召回结果的多样性与覆盖率，Beam Width 可达数百（本文测试为最终 900 条 Beam），且每步可动态变化（如 20 → 300 → 900）。
这个形态和通用LLM Serving框架的优化目标是错位的。业界主流框架（vLLM / SGLang / TRT-LLM）优先优化的都是多用户、长 Decode、单条 Beam 持续生成的对话形态，Beam Search 往往不是一等公民，甚至需要业务方在框架外循环调用才能拼出来。把它们直接搬到 GR 场景，会依次踩中几个坑：
因此我们认为：GR 场景不该由通用框架“顺带支持”，需要一个小而专的引擎，把上下文共享、Beam 状态、动态宽度、受限生成表达成 Runtime 的一等公民。
3.1 设计原则：Request-Centric 状态归属
GR-Inference 的顶层设计以Request（请求）而非独立 Beam 作为状态归属与资源管理的基本单位。该原则贯穿全栈：
调度与资源：生命周期管理、请求准入、KV 资源分配均以 Request 为粒度进行。
计算展开：仅在进入 GPU Decode 阶段后，计算拓扑才按 Batch × Active Beam Width 展开；在此之前，Beam 的差异对上层调度透明。
这一设计分离了逻辑状态树（Beam Search 的分支拓扑）与物理存储单元（KV Cache），使得 Continuous Batching 和内存复用可以在请求级别统一决策，避免以 Beam 为粒度带来的状态碎片化和内存膨胀。
3.2 总体架构
层间交互遵循控制流与数据流分离：Runtime 以上层处理逻辑状态与策略，Backend 以下层处理张量计算；State & Resources Layer 作为中间层，向上屏蔽物理内存布局，向下提供算子可直接消费的索引与指针。
Figure 2. GR-Inference 五层系统架构
状态与资源层
基于 Request-Centric 原则，系统将 KV Cache 显式拆分为三个协作对象，解决“长 Context 共享”与“Beam 分叉隔离”的矛盾：
ContextKV（请求级共享长序列）：保存 Prefill 阶段的长 Context KV，每个请求仅保存一份，其下所有 Beam 共享只读访问，与请求绑定，直到请求结束才释放。
BeamKV（各个Beam私有）：仅保存 Decode 阶段每条 Beam 新产生的短 KV 序列，采用 Step-major 组织，便于按 Decode Step 顺序追加与回溯，Beam 间不共享，确保各分支的生成历史独立。
BeamPath（分叉拓扑与回溯索引）：记录每一步 Beam 的父子关系，用来表示 Beam 的分叉和重排，替代传统方案中“复制整段 ContextKV”的做法。同时在 Decode Attention 中提供 BeamKV 的不规则访问索引，并在输出阶段负责将各 Beam 的 token 序列回溯成最终的 SID 结果。
控制与调度层
Continuous Batching：调度器在合批时保留请求内部的共享关系，对于固定 Beam 宽度请求，通过每个请求的运行时 Step 元数据（Context Length、Decode Step），可以跨 Decode Step 合并至同一批次；对于动态 Beam 宽度请求，按请求实际 Step 与当前 Active Beam Width 合并至同一批次，确保批内计算同构。
Beam Width Policy：提供 fixed、scheduled 与 score-margin 三种 beam 宽度策略。fixed 策略对应 beam 宽度全局固定，scheduled 策略适用每个 Step 不同但请求时固定的情况，score-margin 策略根据每步的分数分布动态收缩宽度：分数高度集中于前若干条 Beam 时，自动把下一步宽度收缩到分数与最优差距在阈值内的 Beam 数（默认只缩不扩），从而在保证质量的同时控制计算量。
执行后端
Bucket-based CUDA Graph：为规避调度开销与 Kernel Launch 延迟，Decode 阶段采用 CUDA Graph 预捕获，按 Batch × Beam Bucket 的固定计算容量预捕获 Graph，覆盖常见形状组合；超出预捕获 Bucket 范围的动态形状，回退至 Eager 执行保证正确性；Graph Key 仅由 Bucket 的固定形状参数决定，不包含请求身份、实际 Context Length 或 Decode Step。因此，同一张 Graph 可在不同请求、不同 Decode Step 之间无感知复用。
Meta-data driven Replay：Graph 绑定服务级共享的 ContextKV Pool 与 BeamKV Pool，请求的Pool Slot 通过运行时元数据暴露给 Kernel；Replay 前只需更新 Context 长度、Pool Slot、Decode Step、有效 Batch 与 Active Beam Width 等小块元数据；由于请求id、实际 Context 长度与 Decode Step 不进入 Graph Key，同一张 Graph 可复用于不同请求与不同 Step 位置，也无需在 Replay 前后复制整块 KV，从而实现“静态图 + 动态数据”的高效执行。
GPU/Capability-based Kernel Selection：执行层将算子接口与后端实现解耦，每种算子操作可被多个后端以不同精度、不同融合粒度实现，由注册表 + 选择策略统一决定，灵活适配不同的GPU架构。
3.3 核心 Kernel OP 优化
GR Decode Attention
现有大语言模型推理框架Beam Search解码阶段，通常将每条 Beam 视为独立序列执行 Attention 计算，忽略了同请求内各 Beam 在上下文层面前缀共享的特性。这导致长上下文（Context）KV Cache 被反复扫描，引发严重的显存带宽冗余，并随 Beam 宽度增加而显著降低GPU计算效率。
针对该问题，本文首先观察到解码阶段 KV Cache 的访存模式存在显著的工作负载异构性：在上下文段（ContextKV），序列长且在同请求所有 Beam 间完全共享，数据局部性良好，呈现连续密集访问特征；在 Beam 分化段（BeamKV），序列短且各 Beam 独立演进，需依据 Beam 历史路径执行非规则 Gather 操作，访存模式离散稀疏。基于上述分布特征，本文提出 GR-Decode-Attention，一种面向层次化拆解的稀疏—密集协同 Attention 架构，通过三级 Kernel 协作流水线实现：
K1 — 密集 Tensor-Core 上下文注意力。 针对 ContextKV 的连续内存布局与高算术密度特征，K1 采用 Tensor Core 执行矩阵乘加（MMA）运算。为避免超长 Context 场景下的 SM 利用率低，引入 Split-KV 并行化策略，沿归约维度对 KV Cache 进行分片，扩展指令级并行与并发线程块数量，从而最大化存储子系统吞吐。
K2 — 稀疏 CUDA-Core 束注意力。 针对 BeamKV 的非合并访存（non-coalesced）与 Beam 路径相关 Gather 语义，K2 采用轻量级 CUDA Core FMA 计算。该设计刻意规避了 Tensor Core 在处理稀疏短序列时面临的刚性分块约束与同步开销，在低计算密度场景下仍保持执行效率。
K3 — 对数求和指数尾段归并。 通过数值稳定的在线 log-sum-exp softmax 归约机制，K3 融合 K1 与 K2 产生的部分注意力输出，直接生成最终概率分布，无需物化中间完整注意力矩阵。
该架构通过算法分解与 GPU 执行层次的深度协同，调和了密集型计算的吞吐需求与稀疏型访存的延迟敏感性，实现了宽 Beam 配置下的可扩展高性能解码。
GPU SID Trie 与 Constrained Top-K
在Item-constrained生成任务中，模型需在每一解码步保持在合法的 SID（SID）路径上。传统实现依赖 CPU 逐 Beam 查询前缀树（Trie）并构建完整词表掩码（Vocabulary Mask），随后在 GPU 上执行全词表范围的 Masked Top-K。该分层处理范式引入了昂贵的host-device同步开销，并产生冗余计算。
GR-Inference 通过 GPU 常驻的压缩稀疏行（CSR）Trie 与单一融合 CUDA Kernel 消除了上述瓶颈。该 Kernel 将三项操作内聚为统一执行路径：（i）各 Beam 在 SID Trie 上的并行前缀遍历；（ii）基于当前 Trie 节点对合法后续 Token 的细粒度 Gather；以及（iii）仅针对约束候选集合的局部 Top-K 提取。
这样既避免了完整词表 Mask，也将 Top-K 的处理范围从整个词表缩小到当前 Trie 节点的合法分支，在确保生成结果严格满足路径约束的同时，维持了极低的解码延迟与较高的推理吞吐。
Figure 3. GPU Trie Constrained Top-K Kernel 示例
3.4 Benchmark
测试方法
SGLang 是开源社区中较完整地把通用 Beam Search 路线推进到可运行状态的尝试，因此本文选择它作为对比对象。主要考察它们在 SID 生成式场景中的端到端能力，而不只是单个 Kernel 的性能。具体的三套SGLang版本如下：
Version 1：feature/beam_search，commit af4e3f42a5a1
Version 2：feature/beam_search_update_0801，commit 9380595f3392
Version 3：lsyin/beam-search-dev，commit 5203fab5eb97（Draft PR #31626）
需要说明的是，在本文测试的版本中，各框架对业务关键能力的支持并不完整：SGLang 的三套实现均不支持动态 Beam，仅 Version 2 支持 Item-constrained 生成。对于大规模搜推系统，动态 Beam 用于在不同 Decode Step 灵活扩展候选空间，Item-constrained 用于保证生成结果始终落在合法 SID 路径上，二者都是线上业务实际需要的能力。GR-Inference 同时支持动态 Beam、Item-constrained、topk_logprob，以及动态 Beam 下的 Decode CUDA Graph，为后续性能对比提供了更完整的业务功能基础。
我们的测试使用 Qwen3-0.6B 的模型，prefill 长度从 100 逐步加到 1000，decode 长度固定为 3，对比的是在 E2E Latency 不超过 100ms 的情况下各框架能达到的最大吞吐。
测试一: 固定 Beam 900
第一组测试将三步 Beam 都固定为 900。GR-Inference 和三套 SGLang 均启用 Decode CUDA Graph，并统一使用普通的全词表 logprob。这是本文 GR-Inference 与 SGLang 之间打分口径最一致的一组性能测试。
Figure 4. 固定 Beam 900 的 SLA 峰值吞吐对比
测试数据显示，GR-Inference 是 SGLang Version 1 的 1.50～2.29 倍，是 Version 2 的 1.66～2.27 倍。对于 Version 3，在输入长度 100～400 的有效测试点中，GR-Inference 是其 2.04～2.53 倍，Version 3 在更长输入下没有满足 SLA 的点。
另外可以发现，随着输入长度增加，GR-Inference 相对 Version 1 和 Version 2 的优势整体扩大。这一趋势与系统设计目标一致：Context 越长，按请求共享 ContextKV、并让专用 Attention Kernel 直接处理共享 Context 的价值越明显。
测试二: 固定 Beam 900 + Item-constrained能力
第二组测试同样将三步 Beam 都固定为 900，比较 GR-Inference 与唯一支持该功能的 SGLang Version 2，两边均启用 Decode Graph。
Figure 5. Item-constrained SID 生成的 SLA 峰值吞吐对比
测试数据显示，GR-Inference 吞吐是 SGLang Version 2 的 3.24～3.63 倍，同样是大幅领先。
4.1 阶段性收益
GR-Inference 是小红书团队和NVIDIA团队联合开发的针对“长 Context + 短 Decode + 大 Beam”生成的专用推理引擎，在固定 Beam 900 的 SLA 口径下，GR-Inference 相对三套 SGLang 社区实现取得 1.5x~2.5x 收益，Item-constrained 场景相对 SGLang 提升 3.2x 以上。
GR-Inference 在小红书搜索场域完成落地，召回阶段新增一路自回归SID召回通道，在耗时达标的情况下，实现了GPU资源的高效利用，最终项目顺利上线，点击率提升 0.03pt，有效点击率提升 0.2%，离线 Recall@1000 召回率提升 5.7%。
4.2 后续方向
KV Cache 内存优化：面向真实 Context 长度分布引入多 Context Bucket 与 Paged ContextKV，让 Decode Attention 原生支持 Page Table，降低变长上下文的显存浪费。
CUDA Graph 覆盖扩大：扩充 Prefill/Decode Graph 的预热形状，并将 Beam Selection（log_softmax + topK）纳入 Graph，进一步缩短热路径。
框架和算子优化：支持 PD 分离、多种并行策略扩展以支持模型 scaling，做更细粒度的 mega kernel融合。
小红书引擎架构部 AI Infra 团队
小红书引擎架构部AI Infra团队（含模型工程方向）致力于支撑搜索、广告、推荐、审核、创作等核心业务的AI化演进，为搜广推、CV及LLM/VLM业务提供端到端的深度学习高性能训练与推理服务。团队与算法紧密协同，共同为大模型在业务场景中的落地效果负责，聚焦大模型分布式训练、强化学习框架、高性能推理引擎、异构硬件编译器等核心技术的研发，主导SOTA训练/推理引擎的架构设计与核心模块开发，推动长序列建模、生成式推荐、Agent等前沿场景在GPU、XPU等异构算力上的规模化落地。我们持续打造先进、高效、易用的AI基础引擎，为小红书AI业务的持续演进提供坚实的底层支撑。
NVIDIA DevTech 团队
Bin Chai（柴斌），NVIDIA加速计算专家，专注于企业级客户的GPU加速优化支持和GPU软件生态研发。目前主要负责GR-Inference推理相关工作。
Jerry Chen（陈乔瑞），NVIDIA加速计算专家，专注于企业级客户的GPU加速优化支持和GPU软件生态研发。目前主要负责Mega Kernel优化相关工作。
【社招】
【ACE顶尖实习生】