vllm-ascend 的 KDA 状态复制 Triton 集成：PR 17301 对齐的预编译 worker 生命周期方案-平芜编程栈

人工智能
大模型
模型推理服务
Ascend
CANN

【免费下载链接】vllm-ascend

Community maintained hardware plugin for vLLM on Huawei Ascend

项目地址：https://gitcode.com/gh_mirrors/vl/vllm-ascend

导读

本文讲解 vllm-ascend 中 Kimi K3 delta attention（KDA）prefill 状态复制（state copy）的 Triton 集成方案。该方案对齐上游 vLLM PR #17301 的数值语义，将原生 gather/clear/scatter 算子替换为启动期预编译、服务期零 JIT 的 Triton kernel，并通过 worker 持有的KDAStateCopyPlan生命周期管理编译产物。读完本文，你将掌握其自动分发规则、四变体预编译设计、数值契约、Dynamo/FakeTensor 边界，以及配套的验证与性能验收方法。

背景：KDA prefill 为什么需要 state copy

Kimi K3 使用 delta attention，其 Ascend 后端实现位于 kimi_kda.py（AscendKimiK3DeltaAttention，标记_requires_kda_state_copy = True）。在 prefill 阶段，层需要完成「初始状态拷贝 → chunk 数学 → 最终状态回写」三步：先从缓存行按索引 gather 出每条请求的初始递归状态，交给run_chunk_kda（见 kda.py，chunk 大小固定为 64），再把 chunk 计算产出的 final state scatter 回缓存行。这一 gather/scatter 环节正是本方案替换的对象。

值得强调的是替换边界：Triton 只接管「缓存行 ↔ packed 状态」之间的复制 kernel，不涉及chunk 数学、recurrent decode（decode 走run_recurrent_kda）、GLM KDA 或其他无关 kernel。原文档明确写出该边界，源码中_run_prefill也只把gather/scatter替换为 plan 调用，中间run_chunk_kda保持不变（kimi_kda.py）。

自动分发：fused plan 与 masked byte-copy fallback

方案无需任何额外配置或环境变量开关，worker 根据「设备类型 + 缓存布局」自动选择实现（见 kda_state_copy_plan.py 中initialize_kda_state_copy的分发逻辑）：

场景	采用的 plan
A5（950）设备、FP32/BF16、正数`[N,H,V,K]`形状、内层 payload 密集、缓存页不重叠	`KDAStateCopyPlan`（fused plan）
A3 / 非 950 设备、连续或 strided 等不支持布局	`StridedKDAFallbackPlan`（掩码 byte-copy 回退）

判定的精确条件在supports_kda_state_copy（kda_state_copy_plan.py）：

设备必须是 NPU，且is_950()为真——即AscendDeviceType.A5（见 device_config.py）；
dtype 必须是torch.float32或torch.bfloat16；
张量必须是 4 维且shape[0] > 0，内层三维[H,V,K]从后往前必须满足「size 为 1 或 stride 等于累积 payload」的密集约束，且stride(0) >= payload（页不重叠）。

回退 plan 的语义与原 PR 的 byte-pointer 复制算法一致：gather 无效索引读出零、跳过无效 scatter 目标，即使连续缓存不受支持也遵循同一契约。回退路径的清零使用 Torch 操作（packed.zero_()与masked_fill_），避免为不支持设备引入第二个惰性 Triton kernel。两个 plan 在 kda_state_copy_plan.py 中实现，byte-copy 用的底层 kernel 是 batch_memcpy.py 中的batch_memcpy_kernel（预编译、服务期不再进入其 JIT）。

完整数据流：从缓存绑定到 prefill 落盘

原文档给出的完整流程是：

worker 绑定缓存：KV cache（含 KDA recurrent state 张量）在 runner 初始化时完成绑定；
选择 fused / masked byte-copy fallback：按上述规则逐层判定；
在一次性样本上 prepare：用 scratch 编译 kernel 变体，不触碰真实缓存内容；
seal 所有 plan：校验编译配置未漂移后，将每个 plan 置为 sealed；
原子发布：所有层都 prepare/seal 成功后，才把 plan 绑定到各层；
prefill 执行：gather 初始状态 → 不变 chunk 数学 → scatter 最终状态。

两个关键设计约束贯穿始终：

准备失败必须中止启动。initialize_kda_state_copy中任何一层的缓存未绑定、布局非法、编译失败都会抛RuntimeError，绝不会静默发布一个半就绪的 worker（kda_state_copy_plan.py）；
V1 与 V2 runner 都必须在缓存绑定后调用初始化。V1 侧见 model_runner_v1.py（注释明确「缓存布局已最终化并绑定，在 warmup forward、图捕获或请求之前于一次性 scratch 上 prepare」），V2 侧见 model_runner.py。Kimi 层在初始化完成前不得服务任何请求，_run_prefill中_ascend_kda_state_copy is None或_kda_state_copy_ready为假都会直接报错。

数值契约：必须严格遵守的语义保证

原文档与 kernel 源码（kda_state_copy_kernel.py 顶部 docstring 与_kda_state_copy_kernel实现）共同定义了以下契约：

Gather 侧

无效索引（越界、为负）与初始状态 flag 为 false 的行，输出为全零；kernel 通过safe_row = tl.where(valid_row, cache_row, 0)保证地址安全，再用 mask 控制实际读写；
flag 缺失（flags is None）时，所有有效行全部 gather；重复的 gather 索引合法，可重复读同一缓存行；
packed 输出完整写入（含无效行写零），不遗留未初始化数据。

Scatter 侧

无效索引被跳过（内核里copy_mask = in_payload & valid_row）；有效目标索引必须唯一——唯一性是调用方前置条件，本方案不做设备到主机的同步扫描；
初始状态 flag 只描述初始态，绝不抑制最终写入：scatter 忽略 flag 取值，只要索引有效就写入。

内存与寻址

缓存页 gap（stride(0)与 payload 之间的空隙）与 layer 的 storage offset 既不被复制也不被覆盖——地址计算基于元素偏移而非字节，data_ptr已内含 storage offset；
地址乘法前先提升到 INT64（tl.program_id(...).to(tl.int64)、cache_row.to(tl.int64)），因此 INT32 索引输入、相隔超过 4 GiB 的缓存页都能正确寻址，测试test_kda_state_copy_address_beyond_four_gib专门验证了stride(0) * element_size() > 2**32的场景；
packed 输出的 dtype 与缓存 dtype 一致；最终状态先转换为缓存 dtype 并 contiguous 化再 scatter（fallback 的scatter中packed.to(dtype=state.dtype).contiguous()，与 PR #17301 一致）；
空选择（selected = 0）只验证元数据，不启动任何复制 kernel（_launch中if not indices.numel(): return）；
缓存与 packed 存储不得重叠；并发写入需要上游协调——这些都是调用方前置条件，而非本方案做的同步主机扫描。

启动准备：每布局恰好四个编译变体

KDAStateCopyPlan.prepare（kda_state_copy_plan.py）对每个绑定的缓存布局编译恰好四个变体：INT32/INT64 索引 × gather/scatter四种组合。设计要点：

所有变体共享固定的缓存元数据（payload、行数、页 stride）与DEFAULT_KDA_BLOCK_SIZE = 8192（kda_state_copy_kernel.py）；
gather 的 flag 固定为对齐的连续 BOOL；
索引值与 selected count 是纯运行时数据，launch 时只重绑 grid（[(numel, tiles, 1)]），不触发 JIT 调用；0 <= selected <= scheduler.max_num_seqs不会撑大变体表；
准备只在一次性行上运行：构造 1 行 packed、零索引、全真 flag 与一块 scratch backing，index=0保证 scratch 永远不需要跨越真实缓存；为保留缓存的 16 字节对齐类别，prepare 还专门计算了 scratch 相对 backing 的偏移（state_alignment处理）；
所有四个 kernel 指针指向 scratch 而非真实缓存，准备结束后torch.npu.synchronize()再释放一次性分配，同时把异步编译/启动错误提前暴露到启动期。

StridedKDAFallbackPlan.prepare单独准备一个 byte-copy 变体，源/目标指针同样是 scratch（32 字节的 uint8 张量），不引用真实缓存指针。

服务期只调用已存储的编译对象（_launch直接通过CompiledKernel.__getitem__绕过 Triton JIT 参数绑定器，且保持与准备期一致的 12 参数顺序）。以下情况在 launch 前一律拒绝（fail-closed）：

plan 未 seal；
缓存布局签名（cache_signature：设备、dtype、shape、stride、data_ptr % 16）与准备期不一致；
selected 超过max_selected（即scheduler.max_num_seqs）；
出现意外的 JIT pre-run hook。

任何请求都无法新增或逐出变体；不同 worker 布局需重新初始化，相同布局可共享已编译表；plan 不保留请求张量、流句柄或存活数据指针。对齐连续的 INT32 输入保持 INT32（_canonical仅在需要转换/layout/对齐时才复制，vector=True时还检查stride(0) == 1），strided/未对齐的元数据只在必要时归一化且保留 dtype——这种有界的归一化避免把任意 stride/对齐类别变成新的编译器签名。当前流由编译后的 launcher 每次调用时解析。

编译器环境契约：启动期快照而非逐调用扫描

_configuration()（kda_state_copy_kernel.py）快照torch.__version__、triton.__version__，以及以TRITON_、ASCEND_、CANN_、TORCH_NPU_、NPU_、LLVM_为前缀的环境变量和LD_LIBRARY_PATH、LD_PRELOAD、PYTHONPATH。

该快照只在prepare/seal时比较，不在每次 gather/scatter 时扫描。启动后改变编译器环境、后端二进制或已加载实现均不受支持，正确做法是重启并重新 prepare。这是相对旧版严格 wrapper「每次调用探测漂移」的刻意改变，但它并不让已存储的编译对象具备 JIT 编译能力，也不免除张量布局校验（seal里仍检查_configuration与pre_run_hooks）。

原文档特别警告：不要重新引入每次复制都枚举os.environ的逻辑，也不要为减少启动变体数而把每个 INT32 索引张量转成 INT64——这两项在先前实现中都被实测为回归。集成测试test_int32_fast_path_does_not_convert_or_scan_environment正是把_configuration替换成必抛异常的函数，验证服务期 gather/scatter 不触碰环境扫描。

Dynamo 与 FakeTensor 边界

Eager 调用路径直接走预编译 launch（plan.gather/plan.scatter内torch.compiler.is_compiling()为假时的分支）。在 Dynamo 追踪期间，worker 初始化绑定的 plan 会发出vllm::kda_state_gather与vllm::kda_state_scatter两个 opaque 自定义算子（通过direct_register_custom_op注册，见 kda_state_copy_plan.py）：

Fake 实现（_kda_state_gather_fake/_kda_state_scatter_fake）只做无指针元数据校验（4 维、shape[0]>0、索引 1 维且 INT32/INT64、同设备、payload 密集、页不重叠、packed shape 对齐），并推导输出分配 / 声明 cache 突变，不访问任何真实指针或 plan；
实际执行时，op 通过当前 forward context 的no_compile_layers解析出 sealed plan（_context_plan），与其他 Ascend 自定义算子的做法一致；
不引入任何额外的全局 plan 注册表或运行时编译；独立的prepare()而无 worker 绑定的 plan 只是 eager/graph API，不是可导出的 layer-context 句柄。

原文档强调：这些并不是_C_ascend::kda_state_copy的 ABI 替换品，它们只是以 worker 所有的生命周期保留 prefill 复制语义。测试覆盖 FakeTensor、torch.compile(backend="eager", fullgraph=True, dynamic=True)与 NPU graph replay，包括 FP16 strided fallback；但这不能证明全模型编译或所有编译后端的正确性。

验证体系：CPU 路由测试 + NPU 集成测试

原文档列出的验证分层如下：

CPU 测试覆盖设备/布局路由、启动去重、全有或全无发布、真实 prefill 分发；
NPU 集成测试覆盖：与独立构建的 PR #17301 native 算子的可选精确 FP32/BF16 对比、INT32/INT64、动态 selected 数、页 gap/offset、无效索引、graph replay、真实 prefill/chunk 数学、FP16 strided byte-copy fallback；
正式服务测试把 KDA JIT 入口与编译器 API 替换为立即失败的函数，回退测试还禁用 byte-copy JIT 入口；套件单独断言对齐 INT32 输入保持身份、服务期永不调用编译器环境扫描，且改变 selected 数与图输入值在恒定变体表下必须成功。

建议将集成套件与独立的 plan-copy 套件一起运行（后者使用测试本地的 worker 所有 plan 与禁编译器 fixture）：

python -m pytest --confcutdir=tests/ut/ops \ tests/ut/ops/test_kda_state_copy_lifecycle.py \ tests/ut/ops/test_kda_state_copy_integration.py \ tests/ut/ops/test_kimi_kda.py python -m pytest --confcutdir=tests/e2e/nightly/single_node/ops/singlecard_ops \ tests/e2e/nightly/single_node/ops/singlecard_ops/test_kda_state_copy_integration.py \ tests/e2e/nightly/single_node/ops/singlecard_ops/test_kda_state_copy_triton.py

源码层面的对应证据：test_kda_state_copy_triton.py保留了 PR #17301 的 17 个参数化用例（gap/offset/无效行、重复索引、4 GiB 以上寻址、空选择、内层 stride 拒绝、graph replay 换索引换 flag）；test_kda_state_copy_integration.py覆盖直接 launch 的 12 参数签名、65+ 个 selected 数只用四个变体、worker 去重与失败不发布、真实 prefill 与 native chunk 对照、以及可选的与 PR #17301 native 算子差分测试；test_kimi_kda.py在无 worker 启动时用SimpleNamespace提供 stub plan（test_kimi_kda.py）。

性能验收：可复现的配对基准

原文档对性能验收提出严格的方法学要求：

使用相同输入与分配边界对比 native 与 Triton 生产调用，计时区间必须包含输出分配、必要元数据工作和最终状态转换；
新旧实现在同一进程内对比并轮换顺序；
分别记录 eager Event、同步 wall time 与 graph replay；eager Event 区间包含 host dispatch 间隙，不是纯 kernel 时间；
历史「预分配算子」基准只是辅助参考，不能替代生产 wrapper 计时；
本地测试不是完整的模型服务、TP/PP、并发、sleep/wake 或生成质量验收；native 对比复用已安装扩展而非干净 Ascend C 构建；非 950 回退是代码路由但需要额外硬件验证；全模型编译/Inductor 验收不通过算子追踪测试成立；no-JIT 保证只覆盖已准备的 state-copy 路径，不涉及无关模型 kernel。

可复现的配对性能检查命令（需要独立构建的 PR #17301 参考算子，缺失时基准直接失败；各独立正确性测试仍为必选）：

python -m tests.e2e.nightly.single_node.ops.singlecard_ops.benchmark_kda_state_copy_triton \ --output /tmp/kda-round-1.json --seed 41 --trials 100

每个 seed 使用全新进程与新的输出文件（main中output.exists()直接报错）。基准实现（benchmark_kda_state_copy_triton.py）包含生产分配、native 谓词消融、相等预分配边界，以及同一 native callable 的两个标签作为测量对照；每轮随机洗牌顺序；输出生命周期结束于计时区间之外；记录全部原始样本与实现哈希（source_sha256）；warmup 后拒绝 Triton 编译器入口（compiler_entries计数）；分别上报 host enqueue、同步 wall、eager Event 与 100 次 graph replay。永远不要把 eager Event 当作纯 kernel 时间，也不要用中位数相减去组合不同轮次/模式的结果。

统一生命周期与旧策略迁移

整个方案只有一条生命周期：KDAStateCopyPlan.prepare在 scratch 上编译四个变体 →seal校验启动配置 → worker 在所有层 plan 就绪后才发布。kernel 文件 [kda_state_copy_kernel.py](https://link.gitcode.com/i/22eb605e4da00b27dc2e04fbf8e41273) 只保留 kernel 与公共元数据校验，不含进程级 launcher 注册表、签名预算、锁、逐调用环境扫描，也不修改 JIT 对象run方法。测试独立实例化同一 plan。

先前「仅 standalone」策略测试按以下方式迁移（原文档表格）：

旧断言	统一生命周期断言
全局空注册表 / 不可逆全局 seal	未 seal 的 plan 拒绝使用；独立 plan 可各自 prepare 与 seal
未知签名、张量身份、动态 grid	检查缓存布局/对齐；同布局重绑与动态 S 复用四个变体
全局 128 签名预算 / block-size 覆盖	取消：有限 worker 布局集 + 固定 8192 block 取代用户指定签名预算
元数据错误 / 空选择	非法布局与元数据在编译/launch 前拒绝；S=0 不启动任何东西
逐调用编译器环境漂移	取消：prepare/seal 检查漂移；服务期不得扫描配置
全局`kernel.run`替换	取消：仅编译 launch 对象；测试中拒绝编译器入口

配套测试分工明确：test_kda_state_copy_lifecycle.py用隔离的 host stub 测试生产 plan；test_kda_state_copy_triton.py保留原始数值断言但改为调用 sealed 生产 plan 而非独立注册表。worker 初始化处的发布逻辑（kda_state_copy_plan.py）按(plan_type, cache_signature)去重共享编译表，再为每层copy出一份带_layer_name的独立绑定，保证多 worker 生命周期可在同一进程共存。

边界与注意事项小结

本方案是 PR #17301 数值语义在 Ascend NPU 上的对齐实现，替换的是状态复制 kernel，不是 chunk 数学或 recurrent decode；
适用前提：V1/V2 runner 均需在 KV cache 绑定后、warmup/图捕获/请求前完成初始化；Kimi 层初始化前不服务请求；
不支持服务期改编译环境、改缓存布局、超 selected 数或意外 JIT hook，出现任一情况都 fail-closed 拒绝而非降级；
非 950 回退路径为代码路由，硬件验证仍待补充；全模型编译验收不在算子追踪测试范围内。