feat: support DeepSeek-V4 prefill context parallelism

[LMxyzptlk]

NPU TORCH 后端 DeepSeek-V4 / V4-MTP Prefill Context Parallelism (PCP) 支持

概述

本 PR 为 NPU TORCH 后端的 DeepSeek-V4 和 DeepSeek-V4-MTP 增加 Prefill Context Parallelism（PCP）支持。

采用正交的 DP × CP × TP 并行拓扑：

tp_size = world_size / dp_size / cp_size

rank =
    dp_rank * (cp_size * tp_size)
    + cp_rank * tp_size
    + tp_rank

Prefill 阶段将 query token 切分到不同 CP rank。每个 Decoder Layer 执行完成后，对 CP-local hidden state 进行 all-gather，并恢复全局 token 顺序后继续执行。Decode 阶段保持非 CP 执行。

主要修改

1. 正交 CP × TP 通信组

NPU TORCH 后端的 TP group 缩小为：

world_size / (dp_size * cp_size)

• 创建独立的 CP group，组内保持 dp_rank 和 tp_rank 不变，仅改变 cp_rank
• 不同 DP replica 的 CP group 相互隔离
• 通过 ParallelArgs 暴露实际生效的 TP size
• ATB 后端保持原有通信组行为不变

2. DeepSeek-V4 Prefill 切分

新增 DeepSeek-V4 CP context，用于计算：

• 每个 CP rank 的本地 query 范围
• local-to-global token row 映射
• all-gather 与全局顺序恢复所需的 permutation
• CP-local query sequence lengths
• CP-local KV attention extents

Prefill 执行流程：

1. Decoder Layer 前按 CP rank 切分输入 token rows
2. 每个 Decoder Layer 在本地 token rows 上执行
3. Layer 执行完成后对 hidden states 进行 CP all-gather
4. 恢复原始全局 token 顺序
5. Decode 阶段不启用 CP 切分

3. Sparse Attention 与 Indexer 元数据

将 CP-local attention read view 与全局 KV write metadata 分离：

• 读取路径（CP-local）：Sparse Attention 和 Indexer 使用 CP-local query/KV sequence lengths
• 写入路径（全局）：slot mapping、block tables、KV cache 写入、compressor metadata、index-cache 写入、start positions 继续使用全局元数据

同时修复 CP-local query block 与 KV window 末尾的对齐关系，覆盖普通 Prefill 和 cached-prefix 场景。

4. DeepSeek-V4-MTP 兼容

• DeepSeek-V4-MTP 使用与主模型一致的 CP-local query/KV 元数据规则
• 对于 speculative token，在写入缓存或进入下一步 MTP 计算前，按以下顺序执行共识：
  1. 在 TP group 内广播
  2. 在正交 CP group 内广播
• 共识覆盖以下三个位置：
  • Target Prefill 生成的首个 token
  • 每一个 Draft step 生成的 token
  • Target Validate 生成的 accepted tokens
• 正交 CP 下，即使使用 greedy sampling 也会执行 token 共识，避免不同 CP shard 缓存不同的 token

该修改用于修复以下错误：

MTP decode target state token mismatch

该错误的原因是不同 CP shard 独立生成了 speculative token，导致 scheduler 输入 token 与非 driver CP shard 缓存的 target state token 不一致。

5. ACL Graph 兼容

• 移除 cp_size > 1 与 enable_graph=true 不能同时启用的全局限制
• CP Prefill 保持 eager 执行
• Decode 阶段 CP 不生效，因此可以继续使用 ACL graph capture/replay
• CP 开启时，spec-verify chunked-prefill batch 显式回退到 eager 模式

6. 测试

新增单元测试，覆盖：

• 每个 token row 只属于一个 CP rank
• 每条序列在单个 rank 内保持连续切分
• q_len < cp_size 时允许出现空 CP segment
• cp_size=1 时保持 identity
• gather/restore permutation 互逆
• 非均匀 sequence length
• cached-prefix KV extent
• CP-local query/KV 对齐关系
• DeepSeek-V4 和 DeepSeek-V4-MTP 的 CP capability 注册

已知限制

当前 DeepSeek-V4 CP 路径不支持：

enable_schedule_overlap=true