perf: overlap MTP final-draft-to-target graph updates.

[pjgao]

目标与范围

本 PR 只优化 NPU MTP speculative decoding 中 final draft 完成到 target verification graph 启动 的边界。

• base：main@f174bf6943872f79c65a5e2051e2da79d49aa500
• head：1a5338a0ffdd80a209eddefe49545a90c2fea02d
• 包含：draft → target 的静态任务提前准备、graph 内输入更新和 MTP3/4/5 TileLang specialization
• 不包含：PR perf: eliminate validate-to-draft bubbles with asynchronous device-side handoff. #1966 的 target validation → next draft 预启动优化
• 不修改：DFlash PD 分离流程、非 NPU 后端和普通非 MTP decode 流程

这样拆分后，本 PR 的性能归因和评审范围只对应一个 pipeline 边界。

修改前后泳道图

下图只描述本 PR 覆盖的 final draft → target verification 边界；不包含 PR #1966 的 target validation → next draft 预启动。

修改前：final draft 完成后串行准备 target replay

sequenceDiagram
    autonumber
    participant H as Host / Scheduler
    participant D as Draft compute stream
    participant P as Prepare stream
    participant T as Target graph stream

    H->>D: 提交 draft step 1...N
    par Draft 异步计算
        D->>D: 执行 draft 1...N-1
        D->>D: 执行 final draft
    and 可提前准备的 Host 工作
        H->>P: 构造部分 validate host metadata
    end
    D-->>H: final draft token ready

    Note over H,T: Draft → Target 空泡开始
    H->>H: cast + stack 全部 draft tokens
    H->>P: 多次 copy_ 更新 position / KV / slot / block table
    H->>H: ATB Setup 并更新 paged-attention tiling
    H->>T: 更新 causal-conv graph task host 参数
    H->>T: 提交 target graph replay
    Note over H,T: Draft → Target 空泡结束

    T->>T: target verification
    T-->>H: target output

Loading

修改前，final token 已经在设备侧生成，但 target graph 仍需等待 Host 完成 token 拼装、persistent metadata 更新、tiling Setup 和 graph-task 更新。

修改后：动态依赖在 Device / Graph 内完成

sequenceDiagram
    autonumber
    participant H as Host / Scheduler
    participant D as Draft compute stream
    participant P as Prepare / event stream
    participant T as Target graph stream

    H->>H: 查询 target capability 与 graph key
    Note right of H: e82a16f63 / 71c706eb6
    H->>P: 构造稳定地址的 validate metadata source
    Note right of P: b84bc2c23 / 32a7fa3b4
    H->>T: 预查 static graph-task variant

    par Final draft 异步计算
        H->>D: 提交 final draft
        D->>D: 生成 final token tensor
    and 与 final token 无关的准备
        H->>P: 准备 positions / KV / slots / linear-state metadata
        H->>P: 预提交 static graph-task ready event
    end

    D-->>P: final token ready / device event
    P-->>T: graph stream 等待 device event
    Note over H,T: 不执行 Host synchronize，只保留真实数据依赖

    T->>T: fused token update
    T->>T: fused metadata update
    Note right of T: fbe81b5de
    T->>T: dynamic paged-attention tiling update
    Note right of T: 32a7fa3b4
    T->>T: target verification graph replay
    T-->>H: target output

Loading

修改后，draft token tensor 直接成为 target graph 的稳定 source；graph 前导节点负责 token/metadata 更新，模型 graph 内更新动态 tiling，stream 间使用 device event 表达依赖，因此不再在 final draft 后串行等待整套 Host 更新。

问题

优化前，最后一个 draft token 已生成后，target verification graph 仍需串行等待 host 完成：

1. verify token 拼装与 dtype 转换；
2. positions、KV lengths、slots、block table、linear-state metadata 更新；
3. expanded paged-attention tiling 更新；
4. ACL graph task 更新与 replay 提交。

这些工作全部位于 final draft → target 的关键路径，形成 host dispatch 空泡。

实现

1. 提前准备 target graph task

在 final draft 执行期间准备与最终 token 无关的 target graph task，并为动态输入保留稳定地址。final draft 完成后只保留真实数据依赖，不再重新构建整套 target 输入。

2. 在 target graph 开头融合动态输入更新

新增三个 TileLang kernel family：

• spec_verify_token_update：写入 base token 和 draft tokens；
• spec_verify_metadata_update：更新 positions、KV lengths、slots、block table、linear state 等 persistent metadata；
• spec_verify_attention_tiling_update：更新 expanded paged-attention 动态 tiling 字段。

三类更新均被 target graph 捕获，避免 replay 前额外的 graph-external D2D、同步 memcpy 或 ATB Setup。

3. 用能力契约替代模型名硬编码

新增 SpeculativeVerifyCapabilities：

• requires_causal_chunked_prefill
• supports_in_graph_input_update

MTP runtime 查询模型能力，不判断 qwen3_5、mimo 等模型字符串；能力在 target model 加载时缓存，decode 热路径不重复虚调用或字符串匹配。

• Qwen3.5 hybrid：声明并验证当前 expanded-attention/linear-state graph 输入布局；
• MiMo：只声明 causal chunked-prefill，不进入当前 graph input update 快路径；
• 其他模型：默认关闭并使用原保守路径，适配时需由模型层显式声明能力。

4. MTP3、MTP4、MTP5 使用同一条快路径

框架控制流不写死 depth == 5。当前已验证 specialization：

speculative depth	verify width	specialization
MTP3	4	w4 / w4_bt64 / w4_bs128
MTP4	5	w5 / w5_bt64 / w5_bs128
MTP5	6	w6 / w6_bt64 / w6_bs128

其他深度、multi-sequence、地址不稳定或布局不满足要求时自动回退原路径。本 PR 不宣称支持任意 speculative depth 或 MTP8。

代码规模

相对 main@f174bf694：

• 31 个文件
• +2584/-139

主要改动集中在 runtime/ACL graph、三个 TileLang kernel family 及 wrapper、模型能力契约和对应测试。

提交拆分

顺序	commit	评审主题	规模
1	e82a16f63	speculative verification 模型能力契约	3 files，+54
2	fbe81b5de	三个 TileLang update kernel family 及 wrapper	8 files，+737
3	b84bc2c23	persistent graph 输入与融合更新	3 files，+396/-39
4	71c706eb6	executor/ACL graph task 提前准备与 replay plumbing	8 files，+465/-7
5	32a7fa3b4	MTP worker 接入、模型适配和 MTP3/4/5 快路径	4 files，+471/-93
6	1a5338a0f	TileLang 与 ACL graph 测试	5 files，+461

latest-main 性能验证

固定合同：Qwen3.5-2B、target TP2、draft body TP1、NPU 12/13、LCCL；ACL graph、schedule overlap、chunked prefill 开启；graph double buffer、prefix cache 关闭；固定随机种子和输入；输出 25 tokens；warmup 8 后执行 20 requests × 3 rounds。每个正式 run 均为 60/60 成功。

depth	main 25-token E2E	PR 25-token E2E	变化	main TPOT	PR TPOT
MTP3	161.54 ms	147.52 ms	-8.68%	3.620 ms	3.098 ms
MTP4	157.87 ms	148.17 ms	-6.14%	3.430 ms	3.069 ms
MTP5	162.10 ms	159.89 ms	-1.36%	3.552 ms	3.518 ms

MTP5 的 main 为 A-B-A 两侧基线均值（162.64 / 161.56 ms）；候选位于中间。MTP3/4 采用 build/profiling inactive、同设备同合同配对。共享主机外部构建期间产生的 runs 已排除，不用于上表。

main 与 PR 在 MTP3/4/5 的控制响应内容均完全一致：

100017124785
100055688997

候选运行日志分别确认：

• MTP3：spec_width=4, fused_token=1, fused_metadata=1
• MTP4：spec_width=5, fused_token=1, fused_metadata=1
• MTP5：spec_width=6, fused_token=1, fused_metadata=1

三种深度都捕获 static MTP graph-task；MTP5 覆盖 accepted tokens 1–6，未回退到旧慢路径。

完整 artifacts：/home/g00510989/xllm/runs/20260726_pr1962_rebase_main_validation。

构建与测试

• latest main 与 rebased PR 二进制均构建通过；
• acl_graph_executor_test：19/19；
• sequence_mtp_bootstrap_test：1/1；
• spec_verify_token_update_wrapper_test：4/4；
• spec_verify_metadata_update_wrapper_test：3/3；
• spec_verify_attention_tiling_update_wrapper_test：2/2；
• 本 PR 定向测试合计：29/29。

本机不支持当前 safetensors mmap 加载路径，因此 baseline 与 PR 验证二进制都对称叠加 /home/g00510989/xllm/weight-loading-fix.patch；该 patch 仅用于本机启动和 benchmark，不属于本 PR，也未提交。CANN 9.0 的本地 xllm_ops OPP 构建兼容 patch 同样只用于构建环境，未进入本 PR。

回退边界与其他框架路径

以下任一条件不满足时使用原实现：

• 模型未声明当前 graph input layout capability；
• 非 NPU，或 ACL graph/feature gate 关闭；
• multi-sequence；
• speculative depth 不在 MTP3–5；
• stable source address、graph key、block-table width或 tensor layout 不匹配。

因此 Kimi、DeepSeek、DFlash 及其他尚未适配模型不会误入 Qwen3.5 fast path；它们的共享 runtime 路径保持原保守行为。

[pjgao]

PR #1962 已刷新到 a74b00f56，本次 force-with-lease 更新基于最新 main@abe92b83f。

本轮主要变化：

• PR 范围收敛为 final draft → target verification graph，不包含 perf: eliminate validate-to-draft bubbles with asynchronous device-side handoff. #1966 的 target → next draft 优化；
• 删除 Qwen3.5/MiMo target model 字符串门控，改为模型声明 SpeculativeVerifyCapabilities；
• 将旧的 handoff kernel/API 统一重命名为 spec_verify_*_update，明确表达 graph 输入更新语义；
• MTP3/4/5 共用一条 runtime 快路径，对应 verify width 4/5/6；
• 其他模型、深度和布局默认走保守 fallback。

当前 head 已通过改动 production translation units 编译、TileLang kernel 生成及 NPU wrapper 单测 9/9。PR 描述已区分可归属于 #1962 的约 148.73 ms E2E 与 #1962+#1966 组合版本的约 145.9 ms，避免错误归因。

当前 PR 保持 Draft；切换 Ready 前会补充精确 head a74b00f56 的 main/MTP3/MTP4/MTP5 同合同正式 A/B。

[pjgao]

按评审主题重新拆分提交并更新远端，当前 head 为 26abdf093：

1. 391a9855f：模型能力契约，3 files / +54
2. 2c5ba940b：TileLang graph update kernels，8 files / +737
3. b73cddbb8：persistent graph inputs，3 files / +396/-39
4. 6edabdbf1：ACL graph task 提前准备，8 files / +465/-7
5. fd80f16b4：MTP runtime 与 MTP3/4/5 接入，4 files / +471/-93
6. 26abdf093：测试，6 files / +464/-3

拆分前 a74b00f56 与拆分后 26abdf093 的 Git tree 都是 d0de856c31eb106821900bca929ba2774e9fd900，最终源码内容字节级一致；本次仅整理提交历史，没有改变实现和既有验证结论。

旧的已发布历史仍保存在本地 backup/pr1962-pre-review-split-a74b00f56，如有问题可以精确恢复。

[pjgao]

已完成 latest-main rebase 与当前 head 复验：

• base: f174bf6943872f79c65a5e2051e2da79d49aa500
• head: 1a5338a0ffdd80a209eddefe49545a90c2fea02d
• rebase 后 PR 状态：MERGEABLE
• 定向测试：29/29 通过
• Qwen3.5-2B TP2、draft TP1、25-token、60 requests：
  • MTP3: 161.54 → 147.52 ms（-8.68%）
  • MTP4: 157.87 → 148.17 ms（-6.14%）
  • MTP5 A-B-A: 162.10 → 159.89 ms（-1.36%）
• MTP3/4/5 分别命中 spec_width=4/5/6，fused_token=1、fused_metadata=1，并捕获 static MTP graph-task。
• main/PR 三种深度的控制响应内容一致。

PR 描述已同步刷新提交拆分、代码规模、latest-main 数据和验证边界。共享主机外部 build active 期间的 runs 已排除；本机 mmap 权重加载 patch 与 CANN 9.0 xllm_ops 兼容 patch 均对称用于验证，未提交到本 PR。