[Ascend] Migrate attention to FIA and unify graph updates

[wanfengcxz]

PR 改动说明

背景

Ascend Attention 之前仍有多条 ATB 实现，包括 prefill、Dense paged decode 和
MLA paged decode。为了支持 DeepSeek-V2 MLA 以及统一 eager/graph 行为，本 PR
将相关路径迁移到 torch-npu FIA/FIA v2，并补齐 MLA paged prefill。

CUDAGraph 同时还维护了两套动态参数更新机制：

1. FIA/FIA v2 配合 torch.npu.NPUGraph.update()；
2. ATB paged attention 配合 graph-task handle、event 和模块级
   GraphParams。

旧 ATB graph 路径通过 _graph_params 在 Attention kernel 和 Graph Runner 之间
传递 Tensor、task handle、event 以及 MLA 标记；capture-size 计算函数还通过
副作用初始化这些资源。这会造成以下问题：

• graph 资源生命周期和 capture-size 配置耦合；
• _graph_params 持有 KV cache 等 Tensor 引用，rollout reset 时容易残留旧资源；
• 相同 max_batches 命中 lru_cache 时，初始化副作用不会再次执行，不利于 RL
  rollout 后的 re-capture；
• GraphParams.is_mla 依赖算子 capture 时写入全局状态；
• 新旧 Attention 路径长期并存，增加维护和精度验证成本。

本 PR 完成 Attention 算子迁移、序列长度参数优化和 Graph 生命周期重构，并将
运行时基线固定到支持 NPUGraph.update() 的 torch/torch-npu 版本。

主要改动

1. Prefill Attention 从 ATB 迁移到 FIA

普通 prefill 使用 torch_npu.npu_fused_infer_attention_score 替换
_npu_flash_attention/_npu_flash_attention_unpad：

• 常规语言模型使用 TND layout；
• 无 mask 的视觉 Attention 根据输入形状支持 TND 或 BSH；
• 有 causal mask 的标准 Attention、GQA 和 MLA 使用固定 split-fuse mask 与
  sparse_mode=3；
• 无 mask 路径使用 sparse_mode=0；
• FIA 输出复制回预分配的 attn_output，保持现有接口和内存复用语义。

当 query/key 的 head dim 大于 value head dim 时，按照 MLA 语义切分：

Q = [q_nope, q_rope]
K = [k_nope, k_rope]

NoPE 部分作为 FIA 主输入，RoPE 部分通过 query_rope/key_rope 独立传入。
prefill 路径会对切分后的 Q/K NoPE 和 RoPE Tensor 做必要的 contiguous，满足
FIA 输入要求。

2. Dense/MLA decode 迁移到 FIA/FIA v2

Dense paged decode 使用：

npu_fused_infer_attention_score
input_layout = TND
sparse_mode = 0
block_table + block_size

MLA paged decode 使用：

npu_fused_infer_attention_score_v2
input_layout = BNSD_NBSD
sparse_mode = 0
query_rope/key_rope
block_table + block_size

MLA decode 会将 Q 拆为 NoPE/RoPE，将 paged key cache 从
[block, block_size, kv_head, dim] 转为 FIA v2 所需的
[block, kv_head, block_size, dim]，并将 FIA v2 输出恢复到原有
attn_output layout。

3. 新增 MLA paged prefill

paged_prefill_attention 新增可选参数 head_size_v，并通过公共 LLM API 传入
vendor 实现。Camb/MACA 同步接收该参数以保持 vendor 接口一致。

Ascend 通过 key_cache.shape[-1] != value_cache.shape[-1] 识别 MLA cache：

• latent K/V 与 RoPE K 分开传给 FIA v2；
• paged KV cache 转为 [block, kv_head, block_size, dim]；
• 使用 TND_NTD layout、sparse_mode=3 和 causal mask；
• 强制 block_table 为 int32；
• 兼容预分配 attn_output 和直接返回 output 两种调用方式。

非 MLA paged prefill 仍使用 FIA，保留 TND、paged block table 和
sparse_mode=3。

4. 移除 Attention 热路径中的 .tolist()

以下 FIA/FIA v2 参数改为直接传递 Tensor：

actual_seq_lengths
actual_seq_lengths_kv
actual_seq_qlen
actual_seq_kvlen

避免每层 Attention 执行时把序列长度同步到 CPU 并构造 Python list。
NPUGraph.update() 的 MLA cpu_update_input 仍按 torch-npu 接口要求使用 list；
它属于 graph replay 更新元信息，不是每层 FIA 算子调用。

5. 固定 torch/torch-npu 运行时版本

新增集中式校验，要求：

torch >= 2.8.0
torch-npu >= 2.8.0.post1
torch 与 torch-npu 使用相同的 major.minor 版本
torch.npu.NPUGraph.update 可用

校验在 Ascend vendor 初始化时执行。配套 LMDeploy 改动也会在 Ascend backend
初始化时显式调用。运行时不满足要求会直接抛出带版本信息的 RuntimeError，
不再静默退回旧 ATB 路径。

依赖范围同步调整为：

torch>=2.8.0,<2.10.0
torch-npu>=2.8.0.post1,<2.10.0
torchvision>=0.23.0,<0.25.0

DICP 版本识别同步增加当前环境使用的 torch 2.9.0，依赖中增加
packaging 用于可靠解析版本字符串。

6. 删除旧 ATB graph-task Attention 路径

删除以下全局状态和辅助函数：

GraphParams
_graph_params
_graph_capture_sizes
set_graph_params()
get_graph_params()
clear_graph_params()
update_attn_params()
update_decode_attention_params()
update_decode_attention_mla_params()

同时删除 graph_task_group_*、graph_task_update_*、
_npu_paged_attention 和 _npu_paged_attention_mla 的旧 graph 路径。

重构后的整体算子关系为：

Prefill          -> npu_fused_infer_attention_score
Dense decode     -> npu_fused_infer_attention_score
MLA decode       -> npu_fused_infer_attention_score_v2
MLA paged prefill -> npu_fused_infer_attention_score_v2
Graph replay     -> torch.npu.NPUGraph.update

7. 显式传递 MLA 模型信息

LMDeploy Ascend backend 在构造 Graph Runner 时根据模型配置判断：

is_mla = model_config.k_head_dim != model_config.v_head_dim

is_mla 通过构造参数传递给 AscendGraphRunner 和
AscendSingleGraphRunner。Graph replay 根据实例级字段选择更新参数：

Dense: actual_seq_lengths_kv = Tensor
MLA:   actual_seq_kvlen = list

Attention kernel 不再通过 GraphParams.is_mla 向 Graph Runner 回传模型类型。

8. 分离 capture-size 策略和图资源生命周期

_get_capture_batch_size_impl(max_batches) 现在是纯函数，只负责返回 capture
sizes，不再初始化任何图资源。

AscendGraphRunner.get_capture_batch_sizes() 会优先使用
CacheConfig.cudagraph_capture_batch_sizes，未显式配置时才使用 Ascend 默认
策略。因此 capture sizes 可以随 CacheConfig 跨 rollout 保存，而 NPUGraph、
output buffer 和 Tensor 引用仍在 reset 时释放。

9. 修正 Graph Runner reset/capture 生命周期

• AscendGraphRunner.reset() 先调用基类 reset，清除 rollout 局部的
  padding_batch_size；
• 释放所有 single graph runner、output buffer、runner map 和 graph pool 引用；
• 不清除 CacheConfig 中的 capture-size 配置；
• capture 使用 try/finally 恢复 AscendGraphRunner.capturing，避免异常后状态
  残留。

AscendGraphRunner.capturing 仍然保留，因为 MoE graph capture 路径仍需要该
标记；Attention 已不再依赖它维护全局状态。

跨仓改动与合入顺序

该重构同时涉及 DLINFER 和 LMDeploy。

DLINFER：

• 提供新的 AscendGraphRunner(..., is_mla=False) 接口；
• 将 prefill、Dense decode、MLA decode 和 MLA paged prefill 统一到 FIA/FIA v2；
• 为 MLA paged prefill 增加 head_size_v vendor 接口；
• 删除 FIA 算子序列长度参数上不必要的 .tolist()；
• 固定 torch/torch-npu 最低版本并增加 torch 2.9 DICP 识别；
• 删除旧 GraphParams/ATB graph-task 路径；
• 增加生命周期和 Attention 精度测试。

LMDeploy：

• 根据 ModelConfig 判断 is_mla 并传给 DLINFER Graph Runner；
• 初始化 Ascend backend 时显式执行版本校验；
• 复用按设备缓存的固定 FIA causal mask，避免每次 Attention 执行时分配。

建议先合入 DLINFER，再合入 LMDeploy 配套修改，或者在两个 PR 中相互添加依赖
链接。LMDeploy 新调用方式依赖 DLINFER Graph Runner 的 is_mla 参数。

兼容性变化

这是一次有意的运行时兼容性收紧：

• 不再支持 torch<2.8.0；
• 不再支持 torch-npu<2.8.0.post1；
• 不再提供旧 ATB graph-task Attention fallback；
• torch/torch-npu major.minor 不一致时直接报错；
• eager 和 graph 模式均使用 FIA/FIA v2 Attention 路径；
• 删除 Ascend310P 带 mask 场景的 _npu_prompt_flash_attention prefill
  fallback；带 mask 的 prefill_attention 当前仅支持 Ascend910 系列。无 mask
  的 FIA 视觉 Attention 分支不受该限制。

当前依赖上限仍保留 <2.10.0，torch 2.10 及更高版本需要验证后再放开。

测试

测试环境：

Ascend A3
torch 2.9.0+cpu
torch-npu 2.9.0

运行前加载：

source /afs-weights/wangqing/deepseekv2_dev/ascend_env.sh

单元和算子测试：

pytest -q \
  tests/test_ascend_cudagraph_lifecycle.py \
  tests/test_ascend_attention_precision.py

结果：

8 passed

覆盖 capture-size 纯函数、reset 后配置保留、标准/Dense/MLA prefill、MLA paged
prefill、MLA FIA v2 graph replay 动态 KV 长度更新和 MLA decode 精度。版本校验
通过当前 torch 2.9.0/torch-npu 2.9.0 环境的 vendor 导入和两项端到端模型测试
覆盖。

端到端回归：

DeepSeek-V2-Lite MLA graph
  tp=2, batch size=32
  graph capture/replay 成功
  32 路请求完成生成
  exit code=0

Qwen3.5-35B Dense graph
  tp=2
  完成 batch size 64 -> 1 的全部 graph warmup/capture
  Dense decode graph replay 成功
  生成 10 token
  exit code=0

静态检查：

Black check: passed
py_compile: passed
git diff --check: passed
旧 GraphParams/ATB Attention 符号扫描: no matches

风险与后续验证

已覆盖：

• 标准 Attention、QK/V head dim 不同的 Dense Attention 和 MLA prefill 精度；
• MLA paged prefill 和 MLA decode 精度；
• DeepSeek MLA 和 Qwen Dense 均完成真实 graph capture/replay；
• 当前 torch 2.9.0/torch-npu 2.9.0 能通过运行时校验；
• reset 后 capture-size 配置保留行为有回归测试；
• Attention Python 路径没有残留旧 GraphParams/ATB 调用。

仍建议在真实 RL 环境中进行多轮：

capture -> rollout -> sleep/reset -> 更新权重 -> wakeup -> re-capture

重点检查每轮 capture sizes、旧权重/KV cache 引用、NPU 显存增长和 re-capture
输出精度。

Reviewer 建议重点关注

1. model_config.k_head_dim != model_config.v_head_dim 是否覆盖当前所有 MLA 模型；
2. MLA prefill/decode 的 NoPE/RoPE 切分、contiguous 和 KV cache layout；
3. sparse_mode=3 与固定 split-fuse causal mask 是否覆盖各类 prefill 输入；
4. FIA/FIA v2 序列长度参数直接传 Tensor 的 torch-npu 接口兼容性；
5. Dense/MLA NPUGraph.update() 参数名和数据类型是否符合 torch-npu 接口；
6. reset() 是否完整释放 rollout 局部资源并保留 capture-size 配置；
7. Ascend310P prefill fallback 删除是否符合当前设备支持策略；
8. DLINFER 与 LMDeploy 两个 PR 的依赖和合入顺序。

本 PR 范围

本 PR 已整合以下连续改动，应作为一个完整功能链进行 review：

1. prefill、Dense decode 和 MLA decode 迁移到 FIA/FIA v2；
2. prefill 使用 sparse_mode=3 和固定 causal mask；
3. 新增 MLA paged prefill 及 head_size_v vendor 接口；
4. 删除 FIA 算子参数上不必要的 .tolist()；
5. 删除旧 ATB Attention 路径；
6. 固定 torch/torch-npu 最低版本并支持 torch 2.9；
7. 删除旧 GraphParams/graph-task capture 更新机制并修正 RL 生命周期。

其中 Attention 迁移依赖 FIA 能力，Graph 路径统一又依赖
torch-npu>=2.8.0.post1 的 NPUGraph.update()，因此版本约束属于该功能链的一
部分，不再作为独立 PR 拆分。

最终精度测试

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ wanfengcxz
❌ Super User

---

Super User seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}