vllm-ascend PR #18961 解读：给 GLM 的 draft 层 attention 换上 CANN FlashAttn

1. PR 档案

2. 问题：draft 层的 attention 借了别人的引擎（改动前）

GLM 的架构（①）里除了 N 层主模型，还有一件 MTP 模块——前两篇讲布局时它一直在图角落里挂着黄框。这次它当主角：推理时 draft 先连猜 k 个 token，主模型一次前向验证（一轮"猜-验"见下图 3 上半部）。draft 层自己也是一层层 Transformer，每层也有 attention——数学与主模型一字不差；差别全在"被调用的方式"：序列极短、每步 k 次连猜、验证淘汰后下一步重来。

进了框架，每层的 attention 都要由框架配引擎（③的 platform.py 管这件事）。主模型 target 顺利选上 SFA（稀疏注意力后端）；draft 层呢——它查的是同一张选路表，而表里它的特征是"非 MLA、非稀疏、非压缩"，落格到通用 attention 后端，与最普通的 dense 层同桌吃饭（Ⓐ）。通用后端为"大而全"设计：每层各自建注意力元数据、长度依赖 CPU 镜像——拒绝采样刚淘汰过 token，CPU 镜像可能还没同步；每个 draft step 的每层还要各自重建 tiling 计划，一步 k 次猜 × N' 层，备料重复 N'×(k+1) 次（Ⓑ）。这些开销全部吃进每 token 时延：上游 #18717 两节点 P/D 实测，mean TPOT 14.90 ms 里相当一块是备料，不是数学。

ⒶⒷ 速查（后文正文与代码注释均用此编号）

3. 核心思路：窄门换引擎，target 不动

问题住在引擎与备料上，数学没病——那就不动数学，换引擎。本 PR 三条思路：一、新后端：attention/fa_v1.py 三件套（元数据从 device 长度建、tiling 计划按几何特征缓存并在一个 draft step 内层间复用、CANN flash_attn 算子出结果后回填 vLLM 的输出缓冲）；二、窄门：只在七重条件全过时启用——A5 硬件、V1 runner、本层是 draft、draft 必须 eager、target 是 glm_moe_dsa、draft 是 DFlash2、FP16/BF16 未量化 KV（另排除五个未实现的特性），任一不满足静默回老路，行为与改前一致；三、target 不动：主模型照旧选 SFA，"缓存的 target 选路不变"有测试钉死。改后的样子见图 4：

4. 主线讲解：从架构全貌到这次 PR

先对齐地基。本章走两条线：线一·模型架构线——从模型架构看到 vllm-ascend 框架，看模型的每一件住哪、谁调它，4.1 先摆这块推理的朴素 PyTorch 写法并定位到框架加的壳（图 5 给出全局定位 + 4.2 目录地图·图 6）；线二·PyTorch 实现线——从十行示例的 PyTorch 实现看到框架实现，看它被谁一步步算出来、每步为什么（4.3 流水线·图 7 → 4.4 全链时序·图 8 → 4.5 逐步收益·图 9）；两条线在 4.6 的问题处交汇，4.7 讲这次 PR（图 10 至图 13），4.8 至 4.10 三张表收拢。先说结论：attention 的数学一行没动；本 PR 改的全是"引擎与备料"——draft 层选哪个后端、元数据怎么建、tiling 计划建几次、调哪个算子。

4.1 先看这块推理的原始写法：朴素 PyTorch

老规矩，先立对比基准——这块推理（draft 层的投机解码一步）用 PyTorch 朴素写出来长什么样。自己动手写带草稿模型的推理就会这么写：单卡、没有任何框架概念（有删节）。draft 层的 attention 数学就在 layer 里，与主模型一字不差：

# 朴素 PyTorch 写法：投机解码的一步（单卡，无框架）
draft_kv = []                                 # draft 层自己的 KV 缓存
for t in range(k):                            # draft 连猜 k 个 token
    x = embed(last_token)
    for layer in draft_layers:                # draft 模型的 N' 层
        h = layer.self_attn(layer.norm(x), cache=draft_kv)
        # attention：QKV 投影 → QKᵀ → softmax → ·V → o_proj
        h = layer.moe(h)                      # 前馈；GLM 这层是 MoE
        x = h + x                             # 残差
    last_token = lm_head(x).argmax()          # 猜出下一个 token

accepted = verify(draft_tokens)               # 主模型一次前向，验证 k 个猜的
# 收下命中的、淘汰没中的，下一轮 decode 循环重来

这次 PR 动的不是这段数学——数学一行没动。动的是 vLLM 框架给这段朴素代码加的壳：进了框架，每层的 attention 要由框架配引擎——选哪个后端（Ⓐ）、配料怎么备（Ⓑ：元数据/长度/tiling 计划）。裸 PyTorch 里根本没有"配引擎"这件事，attention 就是 layer 里那几行；draft 层被调用方式的特殊（每步 k 次、序列短、每步重来），框架的老壳没有为它定制——问题全住在壳上（ⒶⒷ 代码长相见第 2 章）。

这段代码在全局的什么位置？把模型架构到框架的全貌画出来（图 5）：

看懂这张图，再定位本 PR。架构件一个没变、链子一段没短——变的是"引擎层"：draft 层的 attention 在选路处（platform.py）多了道窄门，全过就换到新后端（attention/fa_v1.py）；门槛挂在硬件画像（device/）上。这就是为什么图 5 的 ★ 落在引擎三件上，而不是任何注意力算子上。

4.2 先认目录：与 attention 相关的目录，一个一个认

地图没变，落点变了：本 PR 的三个源码文件分别落在 attention/（新后端 fa_v1.py）、platform.py（选路总闸）与 device/（硬件画像能力位）——一个管"跑什么"，一个管"谁上场"，一个管"哪台机器能上场"。

4.3 实现线：attention 在 vllm-ascend 原本是怎么完成的

架构全貌看完了；现在换一条线——从 PyTorch 实现切进去，看它被怎么实现。拿十行 attention 示例当解剖标本（target 主模型这条路；draft 层的老路见 4.6），左列十行自上而下，右列是框架接手每行的落点：

十行示例的每一行，框架都有对应落点。这套"逐行挂门牌"的完成过程就是 vllm-ascend 平时跑 attention 的样子——本 PR 不改这条链的任何一环，它改的是链子入口处"这条链配不配给 draft 用"的选路判断。

4.4 全链时序：谁调用谁，每一棒先自己干再交棒

时序图末尾那串自转圈（后端九步）是主模型 target 的 SFA 路径；draft 层的老路更简单——通用后端没有九步流水线，但每一棒"先自己备料再开算"的节奏是一样的。本 PR 动的是第一棒之前的选路与备料，不动调用链本身。

4.5 为什么这么实现：逐步收益

还剩最后一个问题：target 这条路每一步为什么这么实现？各自省的是什么？把流水线中列原样搬到最左，向右逐层展开——怎么做（带门牌）→ 为什么这么设计 → 收益点：

逐行看收益，"省"分三种。省显存的：低秩 latent 让 KV Cache 不随头数翻倍；省带宽的：融合算子让中间结果不落显存；省算力的：稀疏注意力只算选中的块。本 PR 省的是第四种——省备料：不算得更快，是"准备算"更快（device 长度一次取齐、tiling 计划一步一建），这是 TPOT 那 26% 的来源。

4.6 问题：draft 层的 attention 借了别人的引擎

两条线在此交汇。先用 4.1 朴素基准的口径说清毛病：框架给这段朴素代码加的壳，没有为 draft 的调用方式定制——draft 每步 k 次连猜、序列极短、验证淘汰后重来，而老壳按"通用层"对待它：选路查同桌的表（Ⓐ），备料每层每步重做（Ⓑ，长度还可能取到拒绝采样前的旧镜像）。症状是端侧的：两节点 P/D 下 mean TPOT 14.90 ms——数学只占其中一块，备料吃掉的是纯开销。从引擎点进这段代码（改前户口）：

4.7 这次 PR：七重条件窄门 + 新后端三件套

改法两句话：选路总闸在查表之前插一道窄门（七重条件全过 → 返回新后端；任一不满足 → 照旧查表，静默回老路）；新后端 fa_v1.py 三件套接住 draft 的调用方式。改后户口：

老路 ↔ PR 做法，逐行对照（左列=老路一步的备料与开算，中列=PR 做法，右列=门牌与为什么）：

窄门里到底查什么？七重条件一棵树（任一"否"都回老路——所以叫窄门）：

新后端三件套各自干了什么（fa_v1.py，关系化拆解）：

4.8 对照总表：老路 vs 这次 PR，逐项对照

4.9 PR 落点总表

4.10 测试（6 个文件：精度 + 选路矩阵 + plan 复用三道锁）

5. 前后代码对照（带注释）

5.1 选路总闸：查表之前插窄门（platform.py）

目录地图（③）里 platform.py 管选路。改前，draft 层走到这就直接查表；改后，先过窄门：

# platform.py · get_attn_backend_cls（改后，有删节）
# PCP keeps its distributed backend; eligible local draft layers use CANN FlashAttn.
if not attn_selector_config.use_pcp and _validate_flash_attn_backend(backend_key, attn_selector_config):
    # vLLM resolves this import path to the backend class.
    return "vllm_ascend.attention.fa_v1.AscendFlashAttentionBackend"   # ★ 新窄门：全过 → 新后端

backend_map = {                                                       # 老表原样保留
    (True, False, False):  "vllm_ascend.attention.mla_v1.AscendMLABackend",
    (False, False, False): "vllm_ascend.attention.attention_v1.AscendAttentionBackend",  # Ⓐ draft 改前落这格
    ...
}

窄门本体 _validate_flash_attn_backend 就是图 13 的七重条件直译（backend_key 三元组、硬件能力位、V1 config、spec 存在、runner_type=="draft"、enforce_eager、target 的 model_type、is_dflash2_draft、dtype 与 KV 精度、五特性排除）——每一条"否"都 return False 回老路。

5.2 新后端核心：备料一次，层间复用（fa_v1.py）

# fa_v1.py · AscendFlashAttentionImpl.forward_impl（有删节）
# 只在用这个后端的 worker 里才 import CANN——别的部署不多背一个依赖
from cann_ops_transformer.ops import flash_attn, flash_attn_metadata

# Ⓑ 反面一：长度只信 device，惰性算一次，存在 metadata 里
if metadata.flash_attn_lengths is None:
    boundaries = metadata.query_start_loc_gpu[: num_reqs + 1].int()
    boundaries = boundaries.clamp(min=0, max=query.shape[0]).cummax(dim=0).values  # dummy 请求可能有重复/负边界
    lengths = metadata.seq_lens_gpu[:num_reqs].int().clamp(min=1)
    metadata.flash_attn_lengths = boundaries, lengths

# Ⓑ 反面二：tiling 计划按几何特征做 key，一个 draft step 内层间复用
plan_key = (self.num_heads, self.num_kv_heads, query.shape[-1], self.sliding_window, layout_kv)
cached = metadata.flash_attn_plans.get(plan_key)
if cached is None:
    attrs = dict(... mask_mode=4 if ... else 3 if metadata.causal else 0,   # 0 无遮罩/3 因果/4 滑窗
                 layout_q="TND", layout_kv=layout_kv, layout_out="TND")
    cached = flash_attn_metadata(...), attrs
    metadata.flash_attn_plans[plan_key] = cached

with attention_transfer_window():          # KV 池传输窗口内调用
    result, _ = flash_attn(query, self.key_cache, self.value_cache,
                           block_table=..., metadata=plan, **attrs)
    output[: query.shape[0]].copy_(result)  # CANN 返回新张量；vLLM 要自己的输出缓冲——回填
    return output

三个细节值得停一下：clamp+cummax——dummy 请求的边界可能重复甚至为负，不做这两步会算出非法窗口；plan_key 五元组——同几何才共享，换头数/换窗口自动分家（"不同几何不共享 plan"有测试）；copy_ 回填——CANN 的返回是新张量，vLLM 的约定是写进它给的缓冲，一行 copy_ 把两个约定接上。

5.3 硬件门槛：能力位三行（device/hardware_profile.py）

# device/hardware_profile.py（+3）
class HardwareCapability(Enum):
    ...
    # CANN device-side FlashAttn metadata with strided paged GQA K/V.
    DEVICE_LENGTH_FLASH_ATTN = auto()      # ★ 新能力位
# A5 的 supports 列表加上它——条件二的"户口"就在这三行里

6. 与我们排查问题的关系

7. 收益、风险与待验证

8. 术语速查

来源：github.com/vllm-project/vllm-ascend PR #18961（开放中，base: releases/v0.30.0，+594/−0，9 文件）· 上游 #18717（性能数字出处）· 本页为技术解读，行号与条件以 PR diff 为准