PR #16977
已合入 main · commit dee131a8c
性能优化 · 4 文件 +94/-67
作者 weiguihua2 · 2026-09-22 合入 · 代码基线 a71b766ce ·
上游对齐 vLLM main 84030bbe ·
涉及 vllm_ascend/attention/mla_v1.py + attention/context_parallel/mla_cp.py
| 术语 | 含义 |
|---|---|
| MLA Multi-head Latent Attention | 多头潜在注意力：DeepSeek 系模型的注意力实现，把 KV 压缩到低秩 latent 向量再展开，注意力输出还要经过一次 V-up 上投影回到每头的 V 维度 |
| DCP Decode Context Parallelism | 解码上下文并行：超长上下文下把 KV cache 按 sequence 维分片到 dcp_size 个 rank（每 rank 只存 1/dcp_size），decode 时各 rank 算部分注意力再合并 |
| attention backend 注意力后端 | vLLM 定义的注意力实现接口；上游语义（怎么切、怎么合并），后端管执行（用什么算子、怎么调度） |
| FIA Fusion Inference Attention | 昇腾 CANN 的融合推理注意力算子，vllm-ascend 的 MLA/稀疏注意力在 NPU 上的实际执行者 |
| LSE Log-Sum-Exp | 每段 partial softmax 的对数和指数；两段注意力结果靠它做数学上正确的合并 |
| all-to-all 全交换集合通信 | 每个 rank 把自己的分片发给所有 rank、同时收齐别人的分片，用于拼回完整历史注意力 |
| stream / event 流 / 事件 | NPU 上任务按 stream 排队执行；event 用来在流之间建立"谁等谁"的依赖 |
| record_stream 张量跨流生命周期声明 | 告诉分配器：这块张量还要在另一条流上用，别急着回收它的显存 |
| workspace 算子工作区 | 算子执行时临时借用的显存；图模式下地址必须固定，所以两次 FIA 共用一个预分配好的 workspace |
| KV cache 键值缓存 | 注意力用的历史 Key/Value 缓存，decode 阶段的主要访存来源 |
上游 vLLM 为 DeepSeek 系模型定义了 MLA 注意力接口（MLAAttentionImpl）和 DCP 的切分语义——把 KV cache 按 sequence 维分片、决定"谁算历史、谁算当前、怎么合并"。vllm-ascend 负责的是执行：
vllm_ascend/attention/mla_v1.py 里的 AscendMLAImpl（mla_v1.py:795）是昇腾上的 MLA 注意力后端，真正干活的是昇腾 CANN 的 FIA 融合算子；
attention/context_parallel/mla_cp.py 里的 AscendMlaDCPImpl（mla_cp.py:232）混入上游的 DCPImplMixin、继承 AscendMLAImpl——既遵守上游 DCP 的接口约定，又复用昇腾的 MLA 执行能力。这就是典型的 Ascend 适配层结构：上游管语义，昇腾管执行。
DCP decode 一步要把注意力拆成两段：HISTORY 段（历史 token 的 KV，本 rank 只有 1/dcp_size 分片，算完要把 partial 结果经 all-to-all 交换拼回完整）和 CURRENT 段（当前 chunk 的 KV 各 rank 都有副本，本地直接算）。两段各自输出 partial 注意力和 LSE，最后用 LSE 合并（数学上等价于对全序列做一次 softmax）。合并后的输出是 batch-major 布局，还要做一次 V-up 投影才能交给下一层。
改造前这条链路能跑、数值也对，但藏着三处浪费：
.float() 转成 FP32——这是两次全量张量拷贝 + 精度转换，而下游合并其实直接支持 BF16/FP16，这次转换纯属保守写法的浪费；transpose → bmm → transpose，两次显存搬运只为凑布局。decode 阶段每秒要跑几十上百步、几十层，这三处每步都发生，累积起来就是可观的时延。
把镜头再拉近到代码层的现状：history 段 FIA 和 current 段 FIA 是先后两次算子调用（_run_dcp_mtp_split_attention_op 只换 attention_kind 参数），两次调用共用一个预分配好的 workspace（图模式下地址必须固定）；all-to-all 走融合算子 torch.ops.vllm.sfa_dcp_a2a_fused，在一条专门的通信流上执行。改造前的流编排是：主流跑 history FIA → 通信流等事件后做 .float() + a2a → 主流继续跑 current FIA → 主流等通信完成事件 → 合并。这个编排里 current FIA 看似与 a2a 并行，但它排在主流队列上，要和合并、本层收尾抢主流；而 .float() 的两次大拷贝反而插在通信流的最前面，白白推迟了通信的发出时刻。
这个 PR 的方案可以概括成一句话：把 dtype 转换删掉、把计算和通信在两条流上的分工对调、把 V-up 的三次算子融成一次。三者互相独立，可以分别验证——这也体现在测试上：回归测试同时覆盖了"流顺序"和"BF16/FP16/FP32 三种 history 输出 dtype"两组场景。
第一刀：dtype 直通。all-to-all 融合算子 sfa_dcp_a2a_fused 与合并算子 fused_sfa_dcp_lse_combine 本身支持 BF16/FP16，中间那次 .float() 转换并不承担任何语义（合并的数学对 dtype 不敏感，精度由测试锁住），删掉后 history 的 partial 输出和 LSE 以原始 dtype 直接进通信，省掉两次全量拷贝。
第二刀：流职责对调（本 PR 的灵魂）。原编排是"通信在侧流、计算在主流"；新编排把 current 段 FIA 整个搬到侧流（attn_stream），主流腾出来做 history 的打包与 all-to-all。为什么这样对调更优？三个理由：① 依赖更干净——current 段的输入（current_q/current_k/attn_mask）与 history 段的输出没有数据依赖，唯一的约束是两次 FIA 共用 workspace、必须保持先后，用一个 record_event / wait_event 就能表达；② 主流归还给通信——集合通信在主流发起是 HCCL 的常规用法，打包+发出一气呵成，不再被计算插队；③ 重叠窗口对等——current FIA 的整个执行时长与主流"打包 + 通信"的时长重叠，decode 的当前 chunk 只有 1 个 token、计算很快，把它从主流队尾挪走后，主流干完通信只需等一个事件就能合并。
跨流三纪律（这是本方案最值得学的部分）：① 顺序靠事件——history_ready = main_stream.record_event() 在 history FIA 之后打点，侧流 wait_event(history_ready) 后才动 workspace；② 生命周期靠 record_stream——主流分配、侧流使用的输入张量要 record_stream(attn_stream)，防止分配器在侧流还在读时就回收；侧流产出、主流消费的结果反向 record_stream(main_stream)；③ 汇合靠 join——主流 wait_event(current_attn_done) 之后才允许 combine 读侧流结果。
第三刀：V-up 融合 + 回退守卫。torch_npu.npu_transpose_batchmatmul 允许用 perm_x1 / perm_y 声明两次置换，把"transpose → bmm → transpose"融成一次算子调用，省两次显存搬运。但该算子对 batch×K 维度有 65536 的上限（这里 batch 维是 num_heads、K 维是 kv_lora_rank），所以代码里留了守卫：范围内走融合算子，超限自动回退 torch.bmm 原路径——"融合 + 回退"是接入昇腾融合算子的标准姿势，值得记住。
与老流程怎么接：不改任何接口与配置（PR 明确 no user-facing change）；DCP 的切分语义仍归上游，本 PR 只重排昇腾侧的执行编排；两次 FIA 共享 workspace 的图捕获逻辑保持不变（注释同步更新了原因表述）。
还有一条值得学的反面决策：这个 PR 没有顺手把 prefill 路径也切到新编排——prefill 的 current 段是整个 chunk、计算量大，与通信重叠的收益模型完全不同，硬套会引入更复杂的事件链却换不到收益。把方案克制在 decode 路径内，是「一次只改一个执行编排」的典型取舍，也让回归测试的爆炸半径可控。
读图要点：主流纵贯全部五步——current FIA 虽然与 a2a 理论上并行，但它排在主流队列上；通信流的第一个任务不是通信本身，而是两次 .float() 拷贝（红），把通信的发出时刻往后推；V-up 的三次算子（红）也在主流上收尾。
对比图3 与图4：主流从"五步全包"变成"通信为主"——current FIA（绿）整段搬到侧流，与主流的打包+通信重叠；.float()（删除线）整个消失；V-up 从三个算子缩成一个（含守卫）。两条流之间的依赖只有两个事件：history_ready 向下保护 workspace 顺序，current_attn_done 向上供主流 join。
整个故事的舞台是 AscendMlaDCPImpl（vllm_ascend/attention/context_parallel/mla_cp.py:232）的 decode 路径。它继承 AscendMLAImpl（复用昇腾 MLA 的投影与 V-up）、混入上游 DCPImplMixin（遵守上游 DCP 接口）。进入 decode 前元数据已经把 sequence 维切成两份：HISTORY 组带 block_table 与 cp_history_seq_len（本 rank 分片上的分页 KV），CURRENT 组只有当前 chunk 的 K/V 加一个 attn_mask（sparse_mode=3 的当前窗口掩码）。两次注意力调用用的是同一个方法 _run_dcp_mtp_split_attention_op（mla_cp.py:391），只换 attention_kind 枚举——这是理解本 PR 的钥匙：两段计算本身是对称的，不对称的是它们的编排位置。
先看一个容易被忽略的前置步骤（本 PR 只更新了它的注释，mla_cp.py:530 附近）：进入 FIA 之前要为算子准备 workspace（算子工作区）。图模式（ACL Graph）下算子输入输出地址必须固定，所以代码按 num_tokens 预分配一块 workspace、取所有候选里最大的那个，两次 FIA 共用同一块。这决定了后面事件顺序的语义：侧流上的 current FIA 必须等主流上的 history FIA 用完 workspace 才能动它——不是数据依赖，是工作区复用的顺序依赖。
第一处：history FIA 照旧在主流（mla_cp.py:578-590）。用 MLASplitAttentionKind.HISTORY 调 FIA，输入是分页的 history K/V 与 block_table，产出 partial 输出和 LSE。这一段没有改动，是后面所有编排的锚点。
第二处：current FIA 搬进侧流（mla_cp.py:592-621）。这是本 PR 的灵魂，三行一个节奏：history_ready = main_stream.record_event() 在 history FIA 之后打点；五个跨流输入张量（current 的 q_nope/q_pe/k_nope/k_pe 和 attn_mask）逐个 record_stream(attn_stream)，向分配器声明"这些张量还要在侧流上用"；进入 with torch.npu.stream(attn_stream) 后先 wait_event(history_ready)（保护 workspace 顺序），再跑 CURRENT FIA，最后 record_event(current_attn_done)。出了 with 块，结果两个张量反向 record_stream(main_stream)——侧流生产、主流消费，生命周期声明一个都不能少。
第三处：主流直接做 all-to-all，去掉 .float()（mla_cp.py:622-630）。sfa_dcp_a2a_fused 现在在主流上执行，直接吃 BF16/FP16 的 history 输出与 LSE（旧版的 .float() 两次全量拷贝被整体删除），defer_combine=True 把合并推迟到本地一步做。随后 main_stream.wait_event(current_attn_done)——join 点，主流在此等侧流的 current FIA 收工。
第四处：合并且不打包 current（mla_cp.py:632-640）+ V-up 融合（mla_v1.py:1000-1010）。fused_sfa_dcp_lse_combine（来自 ops/triton/sfa_cp.py）把 a2a 收齐的历史分片与侧流直接产出的 current 结果做 LSE 合并——旧版要把 current 也打包进合并输入，新版用 local_output/local_lse 参数直读，省一次打包。最后 _v_up_proj_batch_major 收尾：batch-major 输入不再转置，npu_transpose_batchmatmul 用 perm_x1=(1,0,2), perm_y=(1,0,2) 一次完成两个置换加批矩阵乘；num_heads × kv_lora_rank 在 [1, 65536) 内走融合算子，超限回退 torch.bmm（守卫常量在 mla_v1.py:73）。
测试怎么锁住行为：tests/ut/attention/test_mla_cp.py 扩到 27 个用例，新增流顺序断言（侧流结果必须晚于 history FIA）和 BF16/FP16/FP32 三种 history 输出的参数化；精度侧用 DeepSeek-V3.1-Terminus W4A8 TP8/DCP8 跑 100 条 GSM8K，evaluator 精度 95%——证明删掉的 .float() 确实不承担语义。
view → transpose → torch.bmm → transpose → reshape 三次算子两次搬运；新写法用 npu_transpose_batchmatmul 的置换参数一次完成。守卫常量（mla_v1.py:73）注释写明这是"perm_x1=(1,0,2) 下融合算子的 batch×K 排他上限"。对应方案：第三刀（融合 + 回退）。vllm_ascend/attention/mla_v1.py · AscendMLAImpl · L73、L1000-1010（改后）# Exclusive batch * K limit of the fused op with perm_x1=(1, 0, 2).
TRANSPOSE_BMM_MAX_SUPPORTED_DIM = 65536 # ⚡融合算子上限：超出自动回退 torch.bmm
def _v_up_proj_batch_major(self, x: torch.Tensor) -> torch.Tensor:
"""Keep the DCP result batch-major and fuse both BMM permutations."""
x = x.view(-1, self.num_heads, self.kv_lora_rank) # ↳ batch-major 输入，不再先转置
# The operator's batch dimension is num_heads, not the token count.
if 1 <= self.num_heads * self.kv_lora_rank < TRANSPOSE_BMM_MAX_SUPPORTED_DIM:
# ⚡两个置换 + 一次 bmm 融成单算子（→NPU torch_npu 融合算子通道）
x = torch_npu.npu_transpose_batchmatmul(x, self.W_UV, perm_x1=(1, 0, 2), perm_y=(1, 0, 2))
else:
x = torch.bmm(x.transpose(0, 1), self.W_UV).transpose(0, 1) # ↳ 回退路径：与旧实现等价
return x.reshape(-1, self.num_heads * self.v_head_dim)
block_table（分页 KV 定位）与 cp_history_seq_len，sparse_mode=0 全量因果；注意它产出 history_output / history_lse 两个 partial 张量，正是后面 a2a 的输入。对应方案：第二刀的"锚点"侧。vllm_ascend/attention/context_parallel/mla_cp.py · AscendMlaDCPImpl · L578-590 history_output, history_lse = self._run_dcp_mtp_split_attention_op(
q_nope,
q_pe,
history_k_nope,
history_k_pe,
attn_mask=None,
sparse_mode=0, # ↳ 全量因果掩码（历史段是完整序列注意力）
block_table=decode_meta.block_table, # ⚡本 rank 分片 KV 的分页定位
block_size=block_size,
actual_seq_lengths=decode_meta.actual_seq_lengths_q,
actual_seq_lengths_kv=decode_meta.cp_history_seq_len,
attention_kind=MLASplitAttentionKind.HISTORY, # ⚡与 current 段唯一差别：这个枚举
)
record_event 打点、五个输入 record_stream、wait_event 保护 workspace、侧流计算、再 record_event + 结果 record_stream(main_stream)。对应方案：第二刀（流职责对调）。旧版这段位置是"通信流上 .float()+a2a"，新版位置互换。vllm_ascend/attention/context_parallel/mla_cp.py · AscendMlaDCPImpl · L592-621（改后） # Run current-token attention on the side stream while the main
# stream packs and exchanges history. The ready event also orders
# current attention after the history FIA's shared workspace use.
main_stream = torch.npu.current_stream()
attn_stream = _dcp_mtp_comm_stream() # ⚡侧流（复用原通信流对象，职责已对调）
history_ready = main_stream.record_event() # ⚡纪律一：history FIA 完成后打点
for tensor in (current_q_nope, current_q_pe, current_k_nope, current_k_pe, decode_meta.attn_mask):
if tensor is not None:
tensor.record_stream(attn_stream) # ⚡纪律二：跨流输入声明生命周期，防提前回收
with torch.npu.stream(attn_stream):
attn_stream.wait_event(history_ready) # ⚡纪律一：等 workspace 空闲（非数据依赖）
# Current K/V is replicated. Each rank computes its own Q heads
# and contributes the current chunk once during the local merge.
current_output, current_lse = self._run_dcp_mtp_split_attention_op(
current_q_nope,
current_q_pe,
current_k_nope,
current_k_pe,
attn_mask=decode_meta.attn_mask,
sparse_mode=3, # ⚡当前窗口掩码：只算当前 chunk
block_table=None, # ↳ 当前 KV 是副本，无分页表
block_size=0,
actual_seq_lengths=decode_meta.actual_seq_lengths_q,
actual_seq_lengths_kv=decode_meta.actual_seq_lengths_q,
attention_kind=MLASplitAttentionKind.CURRENT,
)
current_attn_done = attn_stream.record_event() # ⚡侧流完工打点
current_output.record_stream(main_stream) # ⚡纪律二：侧流产出、主流消费
current_lse.record_stream(main_stream)
history_output.float(), history_lse.float()；新版主流直连、dtype 直通（痛点①消失）。defer_combine=True 把合并推迟到下一步；wait_event(current_attn_done) 是 join 点；fused_sfa_dcp_lse_combine 用 local_output/local_lse 直读侧流张量，省掉 current 的打包。对应方案：第一刀 + 第二刀的汇合。vllm_ascend/attention/context_parallel/mla_cp.py · AscendMlaDCPImpl · L622-640（改后） history_recv = torch.ops.vllm.sfa_dcp_a2a_fused(
history_output, # ⚡旧版这里是 history_output.float()——两次全量拷贝已删除
history_lse, # ⚡同上，LSE 也保持原始 dtype 直通
self.dcp_size,
1,
self.dcp_group.unique_name if self.dcp_size > 1 else "",
defer_combine=True, # ↳ 合并推迟到本地（下一步直读侧流结果）
)
main_stream.wait_event(current_attn_done) # ⚡纪律三：join——主流等侧流 current FIA
# Reduce all history shards and the replicated current chunk exactly
# once, reading current FIA tensors directly without packing them.
attn_output = fused_sfa_dcp_lse_combine(
history_recv,
self.kv_lora_rank,
scatter_dim=1,
local_output=current_output, # ⚡直接读侧流产出，current 不再打包
local_lse=current_lse,
)
return self._v_up_proj_batch_major(attn_output) # →NPU 走代码块 A 的融合算子
.float() / .to(dtype) 都要问一句"下游真的需要吗"——由合并算子的 dtype 支持面和精度测试回答，而不是靠保守转换；record_stream 声明生命周期、join 点显式等待——三条缺一，错的是难复现的显存回收竞态或脏读；torch_npu 融合算子时查清维度上限，代码里留守卫回退到等价的朴素路径，上限以内吃性能、超限不崩。| 阶段 | 文字字数 | 图数量 | 图视角 | 代码块数 | 是否达标 |
|---|---|---|---|---|---|
| 一、背景 | 628 字（脚本统计中文正文，下限 600） | 2（图1 业务流程 / 图2 当前实现时序） | 问题视角 ✓ | 0（按模板不要求） | 达标 |
| 二、方案 | 894 字（下限 800） | 2（图3 当前编排 / 图4 目标编排） | 设计视角 ✓ | 0 | 达标 |
| 三、实现 | 1058 字（下限 1000） | 1（图5 执行时序，节点=函数+行号） | 执行/代码视角 ✓ | 4（A/B/C/D，均含 文件+函数+行号+说明+行内注释，节点可对应） | 达标 |
字数为脚本按中文正文（不含图/代码块/指引框）实际统计，每张图上方均有蓝底「这张图回答什么问题」标注。