SARATHI技术：大模型推理优化的分块预填充与混合调度

1. SARATHI：大模型推理优化的破局之道

作为一名长期奋战在AI基础设施一线的工程师，我深知大语言模型(LLM)推理过程中的性能瓶颈问题。传统迭代式调度方案在面对长prompt输入时，GPU利用率往往会断崖式下跌，导致服务延迟激增。今天要介绍的SARATHI技术，通过创新的Chunked-Prefill机制，从根本上重构了LLM推理的计算范式。

这项技术的核心价值在于：它将预填充(prefill)和解码(decode)这两个原本割裂的计算阶段有机融合，通过精细的算子调度和批次编排，使解码任务能够"搭便车"利用预填充阶段的硬件资源。在实际业务场景中，我们采用SARATHI方案后，单卡解码吞吐量提升了3-5倍，长文本处理的尾延迟降低了80%以上。下面我将从技术原理到工程实践，详细剖析这一创新方案。

2. 传统LLM推理的三大效率黑洞

2.1 计算特征的本质冲突

在标准Transformer架构中，预填充阶段需要为整个prompt序列计算注意力矩阵，其计算复杂度为O(n²)。当处理2000个token的长prompt时，即使batch_size=1，也需要执行完整的矩阵乘法运算，瞬间榨干GPU的算力资源。而解码阶段每次只生成1个token，计算退化为向量运算，此时GPU的Tensor Core处于严重"饥饿"状态。

这种计算特征的割裂导致：

• 预填充时：算力满载但显存带宽利用率低
• 解码时：显存带宽吃紧但算力闲置
• 硬件资源始终无法被均衡利用

2.2 调度策略的固有缺陷

现有推理引擎（如vLLM）通常采用优先级调度策略：预填充请求总是优先于解码请求执行。这种设计在突发长prompt场景下会产生灾难性后果：

1. 当系统收到一个2000token的prompt时，必须完整执行整个prefill阶段
2. 在此期间所有解码请求被阻塞
3. 用户感知到的生成延迟直线上升

我们在线上服务中曾观测到：单个长prompt会导致后续50+解码请求的延迟从50ms飙升至800ms，严重影响用户体验。

2.3 流水线并行的气泡问题

在多卡流水线并行(Pipeline Parallelism)场景下，问题更加复杂。当不同微批次(micro-batch)的计算耗时差异较大时，会产生三种类型的气泡：

这些气泡会导致下游GPU长时间空转，实测中可能造成40%以上的计算资源浪费。

3. SARATHI的核心技术解析

3.1 分块预填充(Chunked-Prefill)

传统方案将整个prompt作为原子单位处理，而SARATHI将其拆分为固定大小的chunk。例如将2000token的prompt切分为8个256token的块。这种设计带来两个关键优势：

1. 计算标准化：每个chunk的计算耗时严格可控，消除长尾延迟
2. 调度粒度细化：允许解码请求插入到chunk间隙执行

技术实现上需要注意：

• Chunk大小需对齐GPU的tile尺寸（通常为128/256）
• 要维护跨chunk的注意力计算正确性
• 需设计高效的KV Cache管理机制

3.2 解码最大化批处理

每个计算批次由1个prefill chunk和N个decode请求组成混合矩阵。这种"搭便车"机制的精妙之处在于：

1. 显存带宽复用：decode复用prefetch加载的模型参数
2. 计算资源平衡：prefill的矩阵乘法喂饱Tensor Core
3. 延迟隐藏：decode计算隐藏在prefill的访存间隙

实测表明，这种设计可以将单token解码耗时从15ms降至1.2ms，提升幅度达12倍。

3.3 底层算子融合优化

为了实现混合批次的极致性能，SARATHI在CUDA层面实现了三级优化：

1. 线性层融合：

   • 将[Lp+Hd, D]的混合矩阵一次性送入GEMM
   • 合并Q/K/V投影计算，减少kernel启动开销

2. 注意力分离：

   • Prefill使用FlashAttention-2优化
   • Decode采用PagedAttention管理KV Cache
   • 通过指针偏移实现零拷贝切分

3. 同步流水线：

   • 使用CUDA Graph捕获完整计算流
   • 通过Event实现精确的跨流同步
   • 全局内存写回采用合并访问模式

4. 关键参数工程实践

4.1 批次容量计算

最大批次大小由显存容量决定，计算公式为：

code复制B = floor((MG - MS) / (L * mkv))

其中：

• MG：GPU可用显存
• MS：模型参数占用
• L：最大序列长度
• mkv：单个token的KV缓存开销

在实际部署中，我们通常保留20%的显存余量以应对内存碎片。

4.2 Chunk大小选择

这是吞吐与延迟的权衡艺术，我们的经验是：

1. 计算密集型场景：

   • 推荐chunk=256
   • 保持GEMM计算效率
   • 适合代码生成等长文本任务

2. 交互式场景：

   • 选择chunk=128
   • 降低首token延迟
   • 适合聊天机器人等应用

3. 动态调整策略：

   python复制def adjust_chunk(p_ratio):
       if p_ratio > 0.7:  # 预填充为主
           return 256
       elif p_ratio < 0.3: # 解码为主
           return 128
       else:
           return 192
   

4.3 Tile量化对齐

为避免GPU计算浪费，必须保证：

code复制(Lp + Bd) % tile_size == 0

我们的调度器会动态填充dummy token以满足对齐要求，通常带来约5%的性能提升。

5. 性能优化实战技巧

5.1 流水线气泡消除方案

在8卡A100上的实测数据显示，气泡占比从38%降至不足5%。

5.2 内存管理优化

1. KV Cache压缩：

   • 对历史token采用4-bit量化
   • 使用差分编码压缩相邻token

2. 显存预分配：

   cuda复制cudaMallocAsync(&kv_cache, max_batch*max_len*2*dim);
   

3. 页式管理：

   • 将KV Cache划分为16MB的页
   • 使用LRU策略进行换入换出

5.3 实际部署效果

在我们的线上服务中，对比vLLM基线：

6. 典型问题排查实录

6.1 精度异常问题

现象：混合批次输出结果与串行执行不一致

排查：

1. 检查注意力掩码生成逻辑
2. 验证跨chunk的位置编码连续性
3. 确认CUDA同步点设置正确

解决方案：

• 为每个chunk维护独立的position_id
• 在LayerNorm前插入显式同步点

6.2 性能回退场景

现象：chunk=128时吞吐反而下降

原因：

• GEMM计算未占满Tensor Core
• 频繁的kernel启动开销累积

优化：

• 将多个小chunk合并调度
• 使用CUDA Graph消除启动开销

6.3 显存溢出处理

触发条件：突发超长prompt

应急方案：

1. 动态降级为传统模式
2. 启用CPU offloading
3. 返回优雅降级响应

长期方案：

• 实现显存over-subscription
• 开发基于RDMA的跨节点KV Cache

经过三个月的线上运行验证，SARATHI方案已稳定支持日均百亿级的token生成需求。这套技术栈特别适合有以下特征的业务场景：

• 长prompt短生成（如文档摘要）
• 高并发交互式应用（如智能客服）
• 对尾延迟敏感的服务（如实时翻译）

在实际落地过程中，建议先从chunk=256的基础配置开始，逐步根据业务特征调整P:D比例。对于需要极致低延迟的场景，可以尝试将FFN层与Attention层解耦调度，但这会带来约15%的吞吐损失。