feat: DeepSeek-V4 DSA context-parallel on NPU

[LMxyzptlk]

DeepSeek-V4 DSA Context Parallel (Model 1, CP == TP)

背景 (Background)

问题

DeepSeek-V4 使用 DSA (DeepSeek Sparse Attention)：MLA 低秩 + 滑窗 (SWA) + 压缩 KV (compressor, compress_ratio ∈ {1,4,128}) + lightning indexer 稀疏 top-k 选择。在 xLLM 的 NPU (昇腾) 推理路径上，V4 的 DSAttention 前向 没有任何序列切分：q/kv 都在全量 token 上计算，只有 head 走 TP 分片。

长序列 prefill 场景下，单卡需要承载全量 token 的 attention 计算和 KV，导致：

• 单卡显存峰值高。
• 首 token 时延 (TTFT) 长。

参考实现是 vllm-ascend 的 dsa_cp（vllm_ascend/attention/context_parallel/dsa_cp.py），它把 TP 组直接当作 CP 组 (get_tp_group())，token 流按 TP world size 均分，每个 rank 拿全量 head、本地 token 切片，算完再 all-to-all 转回标准 TP 头分片布局。

目标 (M1，prefill-only)

1. 在长序列 prefill 下把 query token 沿序列维切分到 CP 组各 rank，降低单卡 attention 计算量与显存峰值、缩短 TTFT。
2. 数值与单卡/纯 TP 结果在容差内一致。
3. 不改变 feature 关闭时的任何行为：FLAGS_enable_dsa_cp=false 时所有路径与非 CP 基线字节级一致。
4. 兼容 MTP / 投机解码。

非目标

• 不做 decode 阶段 KV 分片 (DCP)。M1 只切 prefill 的 query token，KV 在 CP 组内全量复制。
• 不做 Ring-Attention。
• 不覆盖 CUDA / MLU / DCU（仅 NPU）。

---

设计 (Design)

采用与 vllm-ascend 一致的 CP == TP 拓扑，而不是 xLLM 已有的独立 cp_size 维度：

• CP 组就是 TP 组：cp_size = tp_size，cp_rank = tp_rank，cp_group_ = tp_group_。
• launcher 保持 config cp_size == 1（world = dp * tp），CP 几何完全从 TP 派生，运行期由单个开关 FLAGS_enable_dsa_cp 触发。

为什么是 CP≡TP 而不是独立 CP？ 早期尝试过独立 cp_size 维度，但出现了 cp_size(2) 与 cp_group world_size(4) 的 group 不匹配：token 按 cp_size 切分，而 gather 却在整个 TP 组上重组，导致输出错乱。CP≡TP 让 "切 token 的组" 与 "分 head 的组" 是同一个组，all-to-all 转置的 reshape 维度天然对齐，消除了这类不匹配。

单层前向数据流 (prefill, CP active)

每个 rank 拿到 全量 token（worker 不切分，见下），在层内做以下变换：

输入: hidden_states  [T_full, hidden]   (每个 rank 全量, 因为 KV 需全量)

(1) q 走本地 token 切片:
    q_b 加载为 FULL-head (单卡 group, 复制到每个 rank)
    → 每个 rank 在自己的 local token 切片上算 ALL heads 的 q
    → partial RoPE 用 local_cos / local_sin

(2) kv / compressor / indexer-cache 走全量 token:
    kv_proj / compressor 在全量 hidden 上算 → 写全量 KV cache
    indexer 的 index-key cache 用全量 hidden + 全量 cu-seqlens

(3) indexer top-k 走本地 q:
    select_qli 的 query / weights / top-k 用 local cu-seqlens

(4) sparse_attn_sharedkv:
    cu_seqlens_q     = local_query_start_loc   (本地 q)
    cu_seqlens_ori_kv = local_kv_start_loc     (右对齐, 本地 q 看到的全量 KV extent)
    seqused_kv       = local_seq_lens
    → attn_output [T_local, num_heads_FULL, head_dim]

(5) 逆 RoPE 用 local_cos / local_sin

(6) all_to_all_4D 头/token 转置 (scatter_idx=2, gather_idx=1):
    [T_local, num_heads_full, head_dim] → [T_full, n_local_heads, head_dim]
    (转置前 pad 到 tokens_per_rank, 转置后 slice 回真实 token 数)

(7) o_a_proj (列并行低秩 einsum) + o_b_proj (行并行 all-reduce)
    ← 标准 TP o_proj, 因为 head 已被 all_to_all 重新分片到 TP 布局

关键点：

• q 本地、KV 全量。这是 DSA-CP 的核心：稀疏 attention 让本地 query 只看它因果可见的 KV，但 KV 本身在 CP 组内复制（M1 不分片 KV）。
• local_kv_start_loc ≠ local_query_start_loc：本地 query 的因果可见 KV 是右对齐的（最后一个 rank 看到全量 KV 流），所以 ori_kv 路径必须用 local_kv_start_loc（local_seq_lens 的 cumsum），不能复用本地 query 的 cu-seqlens。
• full-head 加载：q_b_proj 与 attn_sink 在 CP 开启时加载为全 head（复制），使每个 rank 能在本地 token 上算所有 head；decode / 非 CP prefill 仍切回 n_local_heads 的 TP 分片，保持基线路径不变。

切分位置：模型侧 (model-side)

DSA-CP 采用 模型侧 token 切分，而不是 worker 侧：

• worker 保留全量 token 流，DSAttention 层内部切出本地 query 切片。
• 因为 DSA 层内同时需要 "本地 q" 和 "全量 KV"，模型侧切分把 AllGather/切分收敛在层内，避免 worker 侧切分后再反向 gather 的语义割裂。这与 vllm-ascend 完全一致。
• 通过 Platform::supports_dsa_cp_model_partition()（NPU-only）+ FLAGS_enable_dsa_cp + cp_size > 1 gate，让 worker 跳过 worker 侧 head-tail CP 切分。

---

代码实现 (Implementation)

改动分层如下（路径相对仓库根）。

1. 配置与开关

xllm/core/common/global_flags.h

• 新增 DECLARE_bool(enable_dsa_cp)：DSA-CP 总开关（默认 false）。
• 新增 DECLARE_int32(dsa_cp_kv_interleave_size)：KV 交错粒度，对齐 vllm-ascend cp_kv_cache_interleave_size。M1 prefill-only 路径忽略，为 M2 decode KV 分片预留。

xllm/core/framework/config/parallel_config.{h,cpp}

• DEFINE_bool(enable_dsa_cp, false, ...) + DEFINE_int32(dsa_cp_kv_interleave_size, 1, ...)。
• 把 enable_dsa_cp 接入 ParallelConfig（domain object）：加入 option_category()、PROPERTY(bool, enable_dsa_cp)、from_flags()、from_json()、append_config_json()。
• XLLM_CONFIG_ASSIGN_FROM_JSON 会把 config.json 的值写回 FLAGS_enable_dsa_cp，所以命令行 flag 与 config.json 两种起服务方式都能读到该开关。

xllm/core/platform/platform.h

• 新增 supports_dsa_cp_model_partition()（NPU-only，constexpr）：报告平台是否支持 DSA-CP 模型侧切分。运行期是否真正走该路径由调用点额外用 FLAGS_enable_dsa_cp && cp_size > 1 gate。

2. CP 本地元数据（核心新增）

xllm/core/layers/common/dsa_metadata.h

• 新增 struct DSACPMetadata：local_query_start_loc / local_seq_lens / local_kv_start_loc / local_start / local_end / tokens_per_rank / num_tokens_pad / local_cos / local_sin。
• DSAMetadata 新增 cp_enabled / cp_size / cp_rank / cp_metadata，取代原来的 cp_input_dict 占位（占位保留向后兼容）。

xllm/core/layers/common/dsa_metadata_builder.{h,cpp}

• 新增 build_cp_local_metadata(query_start_loc, seq_lens, cp_size, cp_rank, full_cos, full_sin)：
  • 把 flatten 后的 query token 流均分到 cp_size 个 rank，返回当前 rank 的本地视图。对齐 vllm-ascend _build_local_token_metadata。
  • 计算 num_tokens_pad（pad 到 cp_size 倍数）/ tokens_per_rank / local_start / local_end；用全局 cu-seqlens 与本地区间求交得 local_query_start_loc / local_seq_lens；local_kv_start_loc 是 local_seq_lens 的右对齐 cumsum。
  • 纯函数，只依赖普通 tensor，可在 CPU 上单测。cp_size == 1 时返回 identity 视图（本地即全量），调用方走单一代码路径。

3. DSAttention 前向（核心新增）

xllm/core/layers/npu_torch/deepseek_sparse_attention.{h,cpp}

• 新增成员：cp_size_ / cp_rank_ / cp_enabled_ / cp_group_ / cp_full_head_。
• 构造函数：cp_enabled_ = FLAGS_enable_dsa_cp && tp_size > 1；启用时 cp_size_ = tp_size、cp_rank_ = tp_rank_、cp_group_ = tp_group_。cp_full_head_ = cp_enabled_。
• full-head 加载：q_b_proj 在 CP 时用 single_rank_group_（world_size==1）的 ColumnParallelLinear，加载全量（未分片）权重 + W8A8 scale，量化 matmul 路径不变。attn_sink 同理按 cp_full_head_ 决定 full-head 还是 TP 分片（构造 + load_state_dict 两处）。
• forward()：cp_active = cp_enabled_ && cp_prefill 时走本地 q + 全量 KV，逆 RoPE 用 local cos/sin，然后：

  auto a2a = parallel_state::all_to_all_4D(o.unsqueeze(0),
                                           /*scatter_idx=*/2,
                                           /*gather_idx=*/1,
                                           /*async_ops=*/false, cp_group_);

  转置前 constant_pad_nd pad 到 tokens_per_rank，转置后 slice 回真实 token 数，再走标准 TP o_proj。
• decode / 非 CP prefill（full-head build 下）把 q 与 sink 切回 n_local_heads_，保证与基线一致。

4. Indexer full/local 分离

xllm/core/layers/npu_torch/deepseek_v4_indexer.{h,cpp}

• select_qli 新增 compress_x / compress_cu_seqlens 可选参数：定义时，index-cache-update 路径（compress_kv）用全量 token 张量，而 query/weights/top-k 仍用本地 x / actual_seq_lengths_query。对齐 vllm-ascend _update_indexer_cache(full) vs _indexer_select_topk(local)。

5. 模型前向装配

xllm/models/llm/deepseek_v4.h

• 派生 CP 几何：FLAGS_enable_dsa_cp && dp_local_tp_size_ > 1 时 cp_size_ = dp_local_tp_size_、cp_rank_ = rank % dp_local_tp_size_；否则回退到 config cp_size。
• 构建 sparse metadata（c1/c4/c128 三种 compress_ratio）时，用 build_cp_local_metadata 把 cu_seqlens_q / seqused_kv / batch_size / max_seqlen_q 替换成本地几何；cu_seqlens_ori_kv 用 local_kv_start_loc（右对齐）。否则 kernel 会按全量几何索引 q，越过本地行数 → MTE DDR 越界。
• indexer 的 qli_metadata 同样本地化（query_lens / key_lens / qli_max_seqlen_q），否则 top-k 塌缩到 position 0（topk_max==0）导致 attention 退化。
• 新增 aux_hidden_states 输出（服务 MTP）。

6. MTP / 投机解码

xllm/models/llm/deepseek_v4_mtp.h

• MTP 复用同一个 DSAttention 层，在 CP 下也切本地 q，所以 draft 的 precomputed metadata 必须同样本地化。之前该文件缺少主模型有的 CP-local 分支，导致 draft crash（MTE 越界）。此 PR 在 MTP 里补齐了与 deepseek_v4.h 对称的 build_cp_local_metadata 分支。
• 从 TP 派生 CP 几何（与主模型一致）。
• forward 传递 aux_hidden_states。

7. Worker 侧 gate

xllm/core/runtime/worker_impl.cpp & xllm/core/runtime/llm_worker_impl.cpp

• 两处 gate：dsa_cp_model_side = FLAGS_enable_dsa_cp && cp_size > 1 && Platform::supports_dsa_cp_model_partition()。
• worker_impl.cpp：needs_cp_prefill_side 增加 !dsa_cp_model_side，跳过 worker 侧 head-tail 切分 + ATB cp-tensor 准备（那是给 V3.2 ATB 路径用的）。
• llm_worker_impl.cpp：uses_worker_cp_partition 增加 !dsa_cp_model_side，避免走 worker 侧 CP 路径（LmHead all-gather + 3-arg logits），V4 不实现该路径。

8. 单测

tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp（新增）+ CMakeLists.txt

• 纯 CPU 单测 build_cp_local_metadata：
  • DocstringRank0/1/2：三个 rank 的 local_query_start_loc / local_seq_lens / local_kv_start_loc 与手算一致。
  • RanksPartitionAllQueryTokens：各 rank 本地 token 数之和 == 全量。
  • CpSizeOneIsIdentity：cp_size==1 返回 identity。
  • PaddingAndBoundaryCross：padding + 跨 rank 边界请求的 KV 长度修正。

---

使用方式

# 命令行
--enable_dsa_cp=true   # 配合 tp_size > 1 (CP 宽度 = TP 宽度)

# 或 config.json
{ "enable_dsa_cp": true }

开关关闭时（默认），所有路径与非 CP 基线完全一致。

---

测试与验证

• ✅ build_cp_local_metadata CPU 单测（本 PR 含）。
• 数值对拍：enable_dsa_cp 开/关的 prefill 输出在容差内一致（bf16 相对误差 < 1e-2）；covering compress_ratio ∈ {1,4,128}。
• 功能矩阵：prefill / chunked prefill / mixed；MTP 开启；tp_size ∈ {2,4}。

已知限制

• 仅 prefill 切 query token，decode 不分片 KV（各 rank 全量 KV 副本）。

[XuZhang99]

add to ParallelConfig

[LMxyzptlk]

done