[Feature] 单机八卡高性能 KV Cache 管理模块开发

[lvyufeng]

【任务领取】MindNLP：单机八卡高性能 KV Cache 管理模块开发

1. 任务背景

在 LLM（大模型）推理过程中，KV Cache 的大小随序列长度线性增长。传统的静态显存分配方式会导致严重的显存碎片化（Fragmentation），限制了长文本处理能力和系统并发吞吐量。为了让 MindNLP 在 Atlas 800 A2 昇腾算力平台上实现工业级推理性能，本项目需要参考业界 PagedAttention 技术，实现一套针对 MindSpore 框架优化的分布式 KV Cache 管理系统。

---

2. 任务目标

基于 MindSpore 框架，开发一套支持单机八卡环境的 KV Cache 管理组件，能够适配 Llama/Qwen 等主流模型进行长序列（如 32K 及以上）的高效推理。

---

3. 详细任务分解

第一阶段：技术调研与方案设计

• 调研分析： 深入研究 vLLM (PagedAttention) 和 LightLLM 的实现原理。
• 设计文档： 在 MindNLP 仓库提交 Issue，详细描述以下内容：
  • BlockManager 设计： 如何管理物理块（Physical Blocks）的分配、引用计数与回收。
  • 映射机制： 逻辑 Token 序列如何通过 slot_mapping 映射到非连续的物理显存块。
  • 架构图： 绘制单机八卡环境下，显存分布、多流协同及数据流向图。

第二阶段：分布式管理模块开发

• KV 显存池初始化： 预先申请 8 张 NPU 上的静态显存（Global Tensor Pool），格式通常为 (num_blocks, num_heads, block_size, head_dim)。
• 核心类实现：
  • BlockAllocator: 实现物理块的申请、释放逻辑。
  • MetadataEngine: 维护每个 Request 的 block_table，记录请求与物理块的映射关系。
• 多卡同步： 确保在分布式推理模式下，各 Rank 间的 Block 分配状态保持同步（逻辑一致性）。

第三阶段：模型集成与算子适配

• PagedAttention 逻辑实现：
  • 在 MindSpore 中通过 ops.Gather 或 ops.ScatterNdUpdate 等算子实现非连续内存的 KV Cache 读写。
  • 适配 Prefill（预填充，一次写入多个 Block）和 Decode（增量生成，逐个 Slot 写入）阶段。
• 模型改造： 选取 MindNLP 中的 LlamaForCausalLM 或 Qwen2ForCausalLM，将其标准的 Attention 层替换为支持 Paged 管理的自定义层。

第四阶段：性能压测与长序列验证

• 正确性验证： 验证开启 KV Cache 管理后，模型输出的 Logits 与原始静态实现的一致性（误差控制在合理范围内）。
• 长序列测试： 在 Atlas 800 A2 上完成长文本（如 32K）推理，证明系统不会因显存碎片化而提前触发 Out of Memory (OOM)。
• 性能指标报告： 统计并对比不同 Batch Size 下的吞吐量（Tokens/s）和首字延迟（TTFT）。

---

4. 验收标准

1. 设计 Issue： 提交至 MindNLP 仓，包含清晰的设计思路和架构图，并通过社区 Maintainer 评审。
2. 核心代码： 提交至 mindnlp/engine 或相关推理目录下，符合 MindSpore 代码规范。
3. 测试脚本： 提供 test_kv_cache.py，支持在单机八卡环境下一键拉起 Llama/Qwen 的长序列推理测试。
4. 技术报告： 提交包含测试数据、显存占用对比及 Atlas A2 适配经验的 Markdown 总结。

---

5. 技术要求与参考

• 硬件平台： Atlas 800I A2 / Atlas 800T A2 (Ascend NPU)
• 软件环境： MindSpore >= 2.3, Python 3.9+, MindNLP
• 关键算子： 建议优先使用 ops.ScatterNdUpdate 进行缓存更新，利用 ops.Gather 或自定义 PagedAttention 算子加速读取过程。