执行摘要

• 一句话：自适应 DSv4 prefill chunk 规划，吞吐提升 4%
• 推荐动作：建议阅读此 PR，了解如何通过动态 chunk 规划优化 prefill 性能。代码重构（将规划逻辑从 forward 移入 metadata builder）是值得注意的设计模式。

功能与动机

PR 指出原来的固定 chunk 大小（4 个请求）不够高效，尤其对于短序列或稀疏 batch。新方案可以“pour in as much requests as possible for one chunk”，从而节省内存和减少 kernel 启动。

实现拆解

1. 新增 CPU 端字段：在 DeepseekSparseSWAMetadata 中添加 prefill_seq_lens_cpu、prefill_query_lens_cpu 等张量，用于在 CPU 上执行规划，避免 GPU 同步开销。

2. 实现 get_prefill_chunk_plan：该方法遍历所有 prefill 请求，对每个请求计算 compressed 区域长度（compressed_lens_cpu）和 gather 区域长度（gather_lens_cpu）。然后贪心地从当前起点向后合并后续请求，直到预测的 workspace 面积（(chunk_size) * (max_compressed + max_gather)）超过允许的最大面积（基于原始 prefill_chunk_size 和模型配置估算）。合并后的 chunk 信息以 (start, end, N, M) 元组列表返回。

3. 修改 _forward_prefill：在 flashmla.py 中，移除原来的固定分 chunk 循环，改用 chunk_plan 动态确定循环边界。workspace 分配（workspace_manager.get_simultaneous）改为每次循环根据当前 chunk 大小动态分配，从而减少总体 workspace 分配。

4. 添加单元测试：在 test_flashmla_sparse.py 中新增测试用例，构造 5 个短序列请求，验证 get_prefill_chunk_plan 返回的 chunk_plan 正确合并为单个 chunk。

关键文件：

• vllm/v1/attention/backends/mla/sparse_swa.py（模块 稀疏注意力；类别 source；类型 core-logic；符号 get_prefill_chunk_plan）: 核心变更：添加 get_prefill_chunk_plan 方法，新增 metadata 字段，实现自适应 chunk 规划逻辑。
• vllm/models/deepseek_v4/nvidia/flashmla.py（模块 DSv4 模型；类别 source；类型 core-logic；符号 _forward_prefill）: 修改 _forward_prefill 方法，使用 get_prefill_chunk_plan 返回的 plan 替代固定分 chunk 逻辑，并动态分配 workspace。
• tests/kernels/attention/test_flashmla_sparse.py（模块 测试；类别 test；类型 test-coverage；符号 test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences）: 添加单元测试验证自适应 chunk 规划的正确性，覆盖短序列 batch 场景。

关键符号：get_prefill_chunk_plan, _forward_prefill, test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences

评论区精华

在 Review 中，zyongye 建议将 chunk 规划逻辑从 flashmla.py 的 forward 方法移到 swa_metadata builder 中，以实现更清晰的分离并支持三种 layer 类型独立规划（swa_only、c4a、c128a）。yewentao256 接受并实施，最终将 get_prefill_chunk_plan 实现在 DeepseekSparseSWAMetadata 中。该讨论解决了代码组织问题，使关注点分离更佳。

• 将 chunk 规划逻辑移到 swa_metadata 中 (design): yewentao256 接受并实施，最终将 get_prefill_chunk_plan 实现在 DeepseekSparseSWAMetadata 中。

风险与影响

• 风险：核心路径变更：prefill 是模型执行的关键路径，修改可能引入回归。但通过 unit test 和 end-to-end 精度验证（GSM8K）已覆盖。

workspace 估算依赖：max_workspace_area 的计算基于 prefill_chunk_size（固定值）和模型配置（max_model_len、window_size 等）。如果这些配置在运行时变化（例如通过 API 调整 max_num_batched_tokens），可能导致估算不准确。但当前流程中这些值在模型加载后固定，风险较低。

新字段初始化：新增的 CPU 张量字段在 metadata 构建阶段必须正确设置，否则 get_prefill_chunk_plan 中的断言会触发崩溃。这有助于快速发现错误。

• 影响：性能：对使用 DeepSeek V4 模型的 prefill 阶段，吞吐提升约 4%（根据 bench 结果）。对 decode 阶段无影响。

内存：通过减少 chunk 数量，降低 workspace 峰值内存占用，可能使更大的 batch 成为可能。

团队：引入了自适应 chunk 规划模式，可供其他模型或 attention 后端借鉴。代码结构更清晰，将规划逻辑与 forward 分离。

兼容性：完全向后兼容，仅影响 DeepSeek V4 模型的 Sparse FlashMLA 后端。

• 风险标记：核心路径变更, workspace 面积估算依赖配置正确

关联脉络

• 暂无明显关联 PR