#52789 [Perf] Support internal prefill checkpoints for Mamba prefix caching, 9%~25% TTFT improvement
原始 PR · 作者 yewentao256 · 合并时间 2026-08-23 01:32
性能优化
重要性 8.90
洞察度 6.00
Mamba 前缀缓存支持内部 checkpoint,TTFT 提升 9%~25%
值得精读。重点看三处设计:一是调度切分(`_mamba_block_aligned_split`)与缓存分配(`_needs_internal_checkpoint`)如何通过 `MambaSpec` 新字段建立跨层契约,避免硬编码模型名;二是 metadata 层用“block 边界 + FLASHKDA_CHUNK_SIZE 对齐”双条件做安全降级,保证非对齐场景不产生错误 checkpoint;三是 Triton 落盘 kernel 一次覆盖 conv 状态与 recurrent 状态的写法。建议与 #53196、#53212、#52971 以及 FlashKDA 上游 PR 连读,理解一个性能特性从内核到调度器的完整演进链。