# PR #21971 完整报告

- 仓库：`sgl-project/sglang`
- 标题：perf: skip KV cache in FA backend for embedding mode
- 合并时间：2026-04-14 07:27
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/21971

---

# 执行摘要
本 PR 在 FlashAttention 后端中引入 `fa_skip_kv_cache` 标志，在嵌入模式且禁用缓存时跳过 KV 缓存读写，直接使用原始 K/V 张量计算注意力，消除每层约 19µs 的开销，提升推理性能，同时确保不影响其他后端。

# 功能与动机
动机源于嵌入模式下 KV 缓存的不必要操作：当配置 `--chunked-prefill-size -1` 和 `--disable-radix-cache` 时，每个请求仅为单次 prefill，KV 缓存被写入和读取但从未重用，浪费约 19µs 每层（`store_kvcache` ~15µs + `prepare_varlen` ~4µs）。优化旨在消除此开销，提升嵌入模型推理效率。

# 实现拆解
实现集中于 `flashattention_backend.py` 文件，关键改动点：
- **条件标志添加**：在 `__init__` 中添加 `self.fa_skip_kv_cache`，基于 `server_args.is_embedding`、`server_args.chunked_prefill_size == -1` 和 `server_args.disable_radix_cache` 判断。
- **缓存写入跳过**：在 `forward_extend` 中，修改条件 `save_kv_cache and not is_cp_mode and not self.fa_skip_kv_cache` 以跳过 `set_kv_buffer` 调用。
- **注意力计算优化**：新增代码路径使用 `flash_attn_varlen_func` 代替 `flash_attn_with_kvcache`，直接处理原始 K/V 张量，并添加断言确保 FP8 KV 缓存 descaling 不支持。

# 评论区精华
Review 讨论聚焦于正确性保障：
> Qiaolin-Yu: "should we assert k_descale, v_descale, and num_splits are none here? since in previous path, these attributes are passed in"

> jasperjiaguo: "Good catch, I passed in num_splits=self.num_splits, so it uses the split heuristics similarly. For fp8 kv it's not relevant for this path w/o kv cache so just added an assert to not silently skip."

结论是添加断言处理 FP8 支持并传递 `num_splits` 参数，确保优化路径安全。

# 风险与影响
**风险**：
- 条件判断错误可能导致缓存未填充，影响 radix cache 查找。
- FP8 KV 缓存 descaling 不支持，限制未来扩展。
- 仅针对 FlashAttention 后端，需确保逻辑隔离避免兼容性问题。

**影响**：
- 用户：在特定嵌入模式下性能提升，但仅限于正确配置的场景。
- 系统：减少 GPU 内核调用，降低计算开销，对整体吞吐量有积极影响。
- 团队：代码维护成本轻微增加，但优化路径清晰。

# 关联脉络
与历史 PR 关联显示持续的性能优化趋势：
- PR #22517：优化 TRT-LLM attention 后端，类似关注计算效率提升。
- PR #22645：添加环境变量控制缓存淘汰间隔，平衡内存与性能开销。
这些 PR 共同体现了 sglang 仓库在 attention 计算和缓存管理上的持续优化方向。