状态（Status）： Draft
作者（Authors）： @yule100
创建日期（Created）： 2026-09-09
更新日期（Updated）： 2026-09-09
相关 Issue/PR： #45924（PR: refactor context parallel attention to v3 API with autograd-flowed stats）

---

1. 概述

1.1 简介

本提案对 NPU 上下文并行（Context Parallel, CP）注意力路径进行重构：将底层算子从 v1 API（npu_fusion_attention / npu_fusion_attention_grad）升级为 v3 API（npu_fusion_attention_v3 / npu_fusion_attention_grad_v3），并移除模块级全局栈 _step_cache_stack，改为通过 op autograd kwargs 流转 softmax 统计信息（softmax_max / softmax_sum / seed / offset）。该重构使 CP 路径对 torch.compile / 流水线并行（PP）友好，并与原生 PyTorch CP（ring attention）的 merge + chunk + single philox 模式保持一致。

1.2 动机

当前 NPU CP 实现存在以下痛点：

1. 全局栈线程安全隐患：_step_cache_stack 作为模块级 list 保存前向中间状态，后向时弹出。虽然 Python GIL 保证线程安全，但该模式与 recompute forward / 编译路径不兼容，且在多 graph task 场景下存在栈污染风险。
2. v1 API 不利于图模式：npu_fusion_attention（v1）返回 int 类型的 seed / offset / numels，图模式下无法被优化器正确追踪，dropout 路径编译失败。
3. Handler 未走原生分发流程：原实现手动 _unwrap_args + DTensor.from_local，绕过了 unwrap_to_op_info 的 sharding propagation，导致 TP+CP 组合场景下 head_num 修正缺失，且 from_local 内部的 _FromTorchTensor autograd.Function 在 dynamo / compiled_autograd tracing 时会 re-dispatch 无策略的 aten 算子（如 detach_），导致编译中断。
4. sharding 策略与 layer-1 割裂：原 CP 策略使用 _op_strategy_context 硬编码 Shard(2)，仅支持 BNSD，且覆盖了 layer-1 的 TP/DP 策略而非扩展。

1.3 目标

目标：

• 将 CP 注意力算子升级为 v3 API，返回 Tensor seed/offset（无 numels），图模式友好。
• 移除全局栈，softmax 统计信息通过 autograd kwargs 流转，对编译/PP 友好。
• Handler 镜像原生 _sdpa_handler 的 unwrap_to_op_info → propagate → wrap 流程，编译安全。
• 使用 @register_sharding 装饰器扩展 layer-1 策略，保留 DP/TP 并追加 CP（seq shard）。
• 补充 _allow_implicit_replication 开关、kwargs 校验、TP+CP head_num 修正。

非目标：

• 不修改 v1 API 的 CP 路径（v1 保留向后兼容）。
• 不实现 pse / padding_mask / prefix / actual_seq_qlen / actual_seq_kvlen / sink 等 per-sequence / per-logit 输入的 ring-step 切分（当前明确 reject 并 raise NotImplementedError）。
• 不修改 native SDPA（flash / efficient / cudnn attention）的 CP 路径。

---

2. 用例分析

场景一：CP 训练（BNSD, causal/non-causal）

• 功能点：q/k/v 按 seq_dim=2 shard 到各 rank，ring attention 前向合并 softmax stats，后向按 step chunk 传递 merged stats + single philox seed/offset。
• 性能指标：通信量为 (world_size - 1) * 2 次 all_to_all / all_gather（前向后向各 world_size - 1 轮 K/V rotate），与原生 CP 一致。
• DFX：
  • 兼容性：支持 BNSD 布局；BSH / SBH / TND 暂不支持（fail-fast）。
  • 可维护性：无全局状态，handler 结构与原生 _sdpa_handler 对齐。
  • 可测试性：test_context_parallel_attention.py 覆盖 forward/backward 数值正确性、通信计数、bf16 ULP 容差。

场景二：TP + CP 组合（多头切分 + 序列切分）

• 功能点：TP shard head_dim 后 local query head 数减少，需在 CP handler 内重算 head_num。
• DFX：head_num 修正逻辑镜像 layer-1 _npu_fusion_attention_handler（line 903-910）。

场景三：编译路径（torch.compile / fullgraph）

• 功能点：v3 API 返回 Tensor seed/offset，可被图优化器追踪；handler 内不使用 from_local，避免 re-dispatch 无策略 aten 算子。
• 约束：当前编译路径在 PyTorch 2.12 仍受 DTensor dispatch + SDPA tracing 限制，测试中 compiled=True 跳过（与 native 对齐）。

---

3. 方案设计

3.1 总体方案

整体设计分为三层：

┌─────────────────────────────────────────────────┐
│  Layer-1: torch_npu.distributed.tensor._attention  │
│  (register_sharding: Replicate / DP / TP)          │
└──────────────────────┬──────────────────────────┘
                       │ @register_sharding 覆盖 → 调用 layer-1 + append CP
┌──────────────────────▼──────────────────────────┐
│  Layer-2: _context_parallel._npu_attention          │
│  ├─ _npu_fa_v3_cp_strategy (forward)               │
│  ├─ _npu_fa_grad_v3_cp_strategy (backward)         │
│  ├─ _npu_fa_v3_handler (forward DTensor handler)   │
│  └─ _npu_fa_grad_v3_handler (backward handler)     │
└──────────────────────┬──────────────────────────┘
                       │ 调用
┌──────────────────────▼──────────────────────────┐
│  Native: _templated_ring_attention / _backward     │
│  (ring K/V rotate + per-step op callback)          │
└─────────────────────────────────────────────────┘

前向流程：

1. unwrap_to_op_info 解包 DTensor → local args + output_spec（含 sharding propagation）。
2. 提取 causal / scale / keep_prob / input_layout，修正 TP head_num。
3. 校验 BNSD 布局、passthrough 参数、sparse_mode。
4. _make_forward_op(step_caches, pt) 创建 per-step closure，内部调用 npu_fusion_attention_v3，online softmax merge（python 或 ring_update op 两种实现）。
5. _templated_ring_attention 执行 ring K/V rotate，每步调用 closure，累积 merged softmax_max/sum。
6. 返回 v3 6-tuple（attn_output, softmax_max, softmax_sum, softmax_out, seed, offset），通过 wrap(local_results, output_spec) 构建 DTensor。softmax_max/sum 作为 DTensor 输出，autograd 会保存并传递给后向。

后向流程：

1. unwrap_to_op_info 解包，local_kwargs 中包含 autograd 流转的 softmax_max / softmax_sum / attention_in / seed / offset。
2. DTensor stats 解包为 local tensor。
3. 从 softmax_max/sum 重建 logsumexp（NPU fa 无 lse 输出）。
4. _make_backward_op_v2(merged_max, merged_sum, seed, offset, pt) 创建 per-step closure，内部调用 npu_fusion_attention_grad_v3，使用 merged stats + single seed/offset（与 native CP 的 single philox 一致）。
5. _templated_ring_attention_backward 执行 ring backward。
6. 返回 grad_q/k/v，通过 wrap 构建 DTensor。

3.2 技术选型

方案	优点	缺点	选择
全局栈 _step_cache_stack	改动最小	编译/PP 不友好、recompute 栈污染、与 native CP 不一致	✗
threading.local	线程隔离	C++ autograd engine 跨线程执行 backward，失效	✗
autograd kwargs 流转	编译友好、与 native CP 一致、无全局状态	需 handler 走 unwrap_to_op_info 流程	✓

3.3 功能与性能设计

功能实现：

• v3 API 调用：npu_fusion_attention_v3 返回 6-tuple（去除 numels），npu_fusion_attention_grad_v3 接受 Tensor seed/offset（去除 numels 参数）。
• softmax merge 两种实现：_merge_softmax_stats_python（纯 Python，fp32 merge）和 _merge_softmax_stats_with_ring_update（调用 torch_npu.npu_ring_attention_update op，BNSD ↔ SBH 适配）。
• IS_CAUSAL partial merge：_templated_ring_attention 在 causal 场景使用 query.chunk(2, dim=2)[1]，closure 内对后半段 stats 单独 merge 后 cat 回完整序列。

性能影响：

• 通信量不变（ring K/V rotate 次数与原实现一致）。
• 去除全局栈的 clone 开销（原实现每步 merged_max.clone() 仍保留，但不再有栈 pop 时的 detach 开销）。
• v3 API 图模式友好，为后续 torch.compile 支持铺路。

3.4 安全隐私与DFX设计

• 兼容性：v1 API 路径不受影响；layer-1 DP/TP 策略通过 @register_sharding 覆盖式保留（先调用 layer-1 函数获取策略列表，再 append CP 策略）。
• 可维护性：无模块级可变状态；handler 结构与原生 _sdpa_handler 对齐，降低理解成本。
• 可测试性：通信计数（CommDebugMode）、ring_update 调用计数、bf16 ULP 容差均有测试覆盖。
• 可靠性：assert 替换为 raise RuntimeError，生产环境 -O 模式下不丢失校验。

3.5 编程与调用设计

3.5.1 编程模型基本设计

• 开发环境：PyTorch >= 2.12（native CP _templated_ring_attention API），torch_npu 支持 v3 算子。
• 开发约束：仅支持 BNSD 布局、full attention / causal sparse_mode 1/2/3；不支持 pse / padding_mask / prefix / actual_seq_qlen / actual_seq_kvlen / sink（明确 reject）。
• 可验收设计：test/distributed/tensor/test_context_parallel_attention.py 的 TestContextParallelAttention.test_ring_attention_sdpa 覆盖 is_causal × compiled × load_balance × rotater × dtype 子测试矩阵。

3.5.2 接口定义与设计

3.5.2.1 npu_enable_cp_dtensor_dispatcher()

• 接口描述：注册 NPU CP forward/backward DTensor handler 和 sharding 策略，开启 _allow_implicit_replication。
• 接口原型：def npu_enable_cp_dtensor_dispatcher() -> None
• 约束说明：需在 CP 计算前调用；会注册 native CP sharding rules（flash/efficient/cudnn）。

3.5.2.2 npu_disable_cp_dtensor_dispatcher()

• 接口描述：移除 NPU handler，恢复 _allow_implicit_replication 原值，注销 CP sharding rules。
• 接口原型：def npu_disable_cp_dtensor_dispatcher() -> None

---

4. 测试设计

• 单元测试：test_context_parallel_attention.py::TestContextParallelAttention
  • test_ring_attention_sdpa：forward/backward 数值正确性（vs 单卡 SDPA），覆盖 is_causal={True,False}、compiled={False,True}、load_balance={False,True}、rotater={ALL_TO_ALL,ALL_GATHER}、test_forward_only={False,True}、use_context={False,True}、dtype={bfloat16,float32}。
  • 通信计数校验：all_to_all_single 次数 = world_size - 1（forward only）/ world_size * 3 - 2（forward+backward）。
  • ring_update 调用计数：若 npu_ring_attention_update 可用则 > 0。
  • bf16 ULP 容差：_bf16_ulp_diff <= 1。
• 集成测试：TestContextParallelUtilities 覆盖 _is_causal_behavior、_context_parallel_shard head-tail / per-document load balancer。

---

5. 缺点和风险

• Breaking Change：_make_backward_op → _make_backward_op_v2、_unwrap_args / _step_cache_stack / _scaled_dot_product_attention_cp_strategy / _register_npu_cp_sharding_rules 等函数移除。若有外部代码引用这些内部符号，需迁移。风险评估：低（均为模块内部私有符号）。
• 编译路径仍不完整：PyTorch 2.12 的 DTensor dispatch + SDPA tracing 限制导致 compiled=True 暂跳过，与 native CP 一致，非本提案引入的回退。
• 功能限制：仅 BNSD、仅 causal/full sparse_mode；per-seq / per-logit 输入未支持。后续需实现 ring-step 切分。
• 维护成本：softmax merge 有 python / ring_update op 两套实现，需同步维护。

---

6. 现有技术

• 原生 PyTorch CP（torch.distributed.tensor.experimental._context_parallel._attention）：_templated_ring_attention / _templated_ring_attention_backward，使用 merge + chunk + single philox 模式，softmax stats 通过 op 输出流转。本提案直接复用其 ring 框架，仅替换 per-step op closure。
• layer-1 torch_npu DTensor（torch_npu.distributed.tensor._attention）：_npu_fusion_attention_handler 走 unwrap_to_op_info + propagate + wrap 流程。本提案的 handler 镜像该流程。

---

7. 未解决问题

• BSH / SBH / TND 布局的 CP 支持何时实现？
• pse / padding_mask / prefix / actual_seq_qlen / actual_seq_kvlen / sink 的 ring-step 切分方案？
• torch.compile fullgraph 路径何时可用（依赖上游 DTensor dispatch + SDPA tracing 修复）？

---

附录

• 参考资料：
  • PR #45924: https://gitcode.com/Ascend/pytorch/merge_requests/45924
  • 原生 CP: torch.distributed.tensor.experimental._context_parallel._attention._sdpa_handler
  • layer-1: torch_npu.distributed.tensor._attention._npu_fusion_attention_handler
• 术语表：
  • CP: Context Parallel（上下文并行）
  • ring attention: 环形注意力，K/V 在 ring 上 rotate 逐块计算
  • softmax stats merge: 在线 softmax 的 max/sum 跨块合并
  • philox seed/offset: dropout 的 RNG 状态，single philox 指全程使用同一 seed/offset