Prhub

#35676 [NPU] DeepSeek-V4 adapt sgl-kernel-npu ops (compressor/sparse-attn/sparse-attn-metadata)

原始 PR 作者 unclezhou486 合并时间 2026-08-25 15:13 文件变更 8 提交数 15 评论 12 代码增减 +93 / -59

执行摘要

NPU 上 DSV4 稀疏注意力适配新算子,metadata 移至 host 计算

PR body 明确指出两个动机:一是为 NPU 平台适配 DSV4 稀疏注意力,对接 sgl-kernel-npu 提供的三个算子(compressor PR#689、sparse_attn_sharedkv PR#708、sparse_attn_sharedkv_metadata_host PR#699);二是把 sparse_attn_sharedkv_metadata 从设备侧搬到 host 侧计算——原先设备张量需要 D2H 同步,'eliminating the D2H sync that drains/stalls the stream and hurts overlapped scheduling'。此外,统一 op 命名空间让 draft 与 normal 路径使用相同 kernel,减少双路径维护成本与行为分叉风险。

值得精读:host 侧 metadata 计算与 CUDA graph 捕获/回放下 CPU 镜像的构建策略是核心设计决策,对理解多硬件(NPU/XPU/AMD)后端适配很有价值;randgun 对简化条件的坚持与 vstone-w 对算子的澄清也值得留意。建议阅读时重点关注 CPU 镜像一致性与 sgl-kernel-npu 版本绑定,若要在生产环境启用,建议先补一次 PR 前后性能对比。

讨论亮点

randgun 在 extra_ops_loader.py 上建议直接删除 initialize_dspark_sparse_attn_ops 及其调用、用注释说明 extra_ops_loader 的用法(已采纳,函数与调用点均已删除,TorchOpLoader 补上详细 docstring);在 allocation.py 上建议去掉 uses_safe_last_loc 中间变量,直接写作 if (_is_hip or _is_npu) and uses_triton_dispatch(head 版本已按此落地)。AndyLi429 询问 host 新算子与 npu_sparse_attn_sharedkv_metadata 是否同一个、是否重命名,vstone-w 澄清两者不同、不能互换;对于测试字段新增,vstone-w 解释 actual_seq_lengths_q_pa_cpu 是 host metadata op 的入参(只接受 host tensor)。ping1jing2 表示 LGTM 但建议补充 PR 前后性能对比,该建议在合入时未落实。

实现拆解

  1. 统一 DSV4 算子命名空间ascend_dsv4_backend.pyforward_compresstorch.ops.custom.compressor 改为 torch.ops.npu.compressor,稀疏注意力相关调用一并替换为 torch.ops.npu.*dspark_worker_v2.pyinit_attention_backends 删除 NPU 分支下的 initialize_dspark_sparse_attn_ops 调用,使 draft 与 normal 路径共用同一套 kernel。
  2. host 侧 metadata 计算ascend_backend.pyForwardMetadata 新增 actual_seq_lengths_q_pa_cpu 字段;ascend_dsv4_backend.pyinit_forward_metadata 中按 prefill(cumsum 前缀和)、decode(0..B 等差)、draft_extend(按 n_draft 步长)三种模式构建 CPU 镜像;_init_dsv4_graph_metadata 在 CUDA graph 捕获期一次构建镜像;_refresh_graph_seq_metadata 新增 else 分支用 final_seq_lens_cpu 维护 seq_lens_cpu_int 兜底,同时修复 _use_host_sparse_metadata 开关关闭时的 fallback;_kernel_metadata_from_parts 的 common 参数中移除 cu_seqlens_qseqused_kv,改由 host op 直接读 CPU 镜像。
  3. 移除冗余算子加载extra_ops_loader.py 删除 initialize_dspark_sparse_attn_ops_C_ascendnpu_sparse_attn_sharedkv_metadata 的必需校验,并为 TorchOpLoader 补充详细 docstring,说明环境变量定位、依赖预加载等用法;dspark_worker_v2.py 不再单独触发加载。
  4. 分配路径与配置读取修复allocation.pyget_last_loc 将安全分支条件从 _is_hip 扩展为 (_is_hip or _is_npu) and uses_triton_dispatch,NPU DSV4 分配路径改走 int32-safe 的 get_last_loc_triton_safedsv4_memory_pool.py 适配 compressor 的 FP32 state-cache 约束;xpu_backend.py 将确定性推理开关的读取从 model_runner.server_args 改为 get_exec().deterministic,符合 config bag 读取规范(配合 CI 的 config ratchet 检查)。
  5. 测试与精度验证test_npu_ascend_backend.pytest_field_names 补充 actual_seq_lengths_q_pa_cpu 断言;PR body 给出 aime26 精度 DSPARK ON/OFF 均为 0.9667、serving benchmark 总吞吐 26686.16 tok/s、TPOT 中位数 29.19 ms,但没有提供 PR 前的基线对比。
文件 模块 状态 重要度
python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py 注意力后端 modified 7.02
python/sglang/srt/hardware_backend/npu/extra_ops_loader.py 算子加载 modified 6.95
python/sglang/srt/mem_cache/allocation.py 显存分配 modified 5.63
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py 投机解码 modified 5.35
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py 注意力后端 modified 4.79
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py 显存池 modified 4.5
python/sglang/srt/layers/attention/xpu_backend.py XPU 后端 modified 4.39
test/registered/unit/npu/attention/test_npu_ascend_backend.py 单元测试 modified 2.91

关键符号

initialize_dspark_sparse_attn_ops get_last_loc get_last_loc_triton_safe forward_compress init_forward_metadata _kernel_metadata_from_parts _refresh_graph_seq_metadata _init_dsv4_graph_metadata init_attention_backends

关键源码片段

python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py core-logic

DSV4 NPU 注意力后端核心,完成 host 侧 metadata 计算(CPU 镜像维护)与 torch.ops.npu 命名空间统一,改动量最大、风险最集中。

# 关键实现单元:为 host metadata 算子维护 CPU int32 镜像输入。
# 背景:NPU 的 torch.ops.npu.sparse_attn_sharedkv_metadata_host 只接受
# CPU int32 张量;若在运行时对 device 张量做 .cpu(),会触发 D2H 同步,
# drain 当前 stream 并打断 overlapped 调度。因此 forward metadata 里
# 新增 actual_seq_lengths_q_pa_cpu 镜像,按 forward 模式构建一次后,
# 直接供 _kernel_metadata_from_parts 组装 kernel 参数时取用。# prefill:每个请求的 KV 边界 = cumsum(seq_lens),前面补 0 得到 cu_seqlens
fm.actual_seq_lengths_q_pa_cpu = torch.cat(
    [torch.zeros(1, dtype=torch.int32), torch.cumsum(seq_lens_cpu, dim=0).int()],
    dim=0,
)# decode:每个请求只产出 1 个 token,边界即 0..B 的等差数列
fm.actual_seq_lengths_q_pa_cpu = torch.arange(0, B + 1, dtype=torch.int32)# graph capture:q_pa 按 graph shape 恒定、replay 从不改写,所以捕获期
# 一次构建 CPU 镜像即可;kv 侧长度则从 CPU 源直接 clamp 生成,同样避开 D2H
metadata.actual_seq_lengths_q_pa_cpu = torch.arange(
    0, bs * tokens_per_req + tokens_per_req, tokens_per_req, dtype=torch.int32
)
fm.seq_lens_cpu_int = ctx.final_seq_lens_cpu[:ctx.bs].int().clamp(min=1)
python/sglang/srt/mem_cache/allocation.py core-logic

将 NPU 纳入 get_last_loc_triton_safe 安全路径,修复 NPU DSV4 分配路径可能出现的 int32->int64 越界问题,影响所有 NPU 平台。

def get_last_loc(req_to_token, req_pool_indices_tensor, prefix_lens_tensor):
    prefill_backend, decode_backend = attention_backends()
    uses_triton_dispatch = prefill_backend not in ("ascend", "torch_native") and \
        decode_backend not in ("ascend", "torch_native")
​
    # HIP 与 NPU DSV4:旧 get_last_loc_triton 内核会写出 int32 -> int64
    # 混合宽度存储,Triton 在 HIP 上会错误编译出越界的 last_loc,NPU DSV4
    # 的等价分配路径也可能 fault。统一走 int32-safe 变体,在 launch 之后
    # 才提升为 int64;其余硬件 / 后端保持原分发逻辑。
    if (_is_hip or _is_npu) and uses_triton_dispatch:
        return get_last_loc_triton_safe(
            req_to_token, req_pool_indices_tensor, prefix_lens_tensor
        )
​
    impl = get_last_loc_triton if uses_triton_dispatch else get_last_loc_torch
    return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor)

评论区精华

删除 initialize_dspark_sparse_attn_ops 设计

randgun:Delete initialize_dspark_sparse_attn_ops func and the code where call it, writes some comments to introduce how to use extra_ops_loader

结论:已删除该函数及其调用点,并扩充 TorchOpLoader 的 docstring 说明用法 · 已解决

get_last_loc 安全路径的条件简化 设计

randgun:Do not need uses_safe_last_loc , change to "if (_is_hip or _is_npu) and uses_triton_dispatch"

结论:采纳建议,head 版本直接使用该条件,去掉中间变量 · 已解决

host metadata 算子与 npu_sparse_attn_sharedkv_metadata 的关系 question

AndyLi429:is it same to npu_sparse_attn_sharedkv_metadata? we rename this kernel?

结论:vstone-w:not same,this is used only for dspark, and can not be replaced by each other · answered

新增 actual_seq_lengths_q_pa_cpu 字段断言的必要性 测试

AndyLi429:why add this? do you test NPU CI and UT?

结论:vstone-w:param for sparse_attn_sharedkv_metadata_host, which only accept host tensor · answered

风险与影响

1) 核心路径一致性风险:ascend_dsv4_backend.py 把 metadata 输入改为 host 侧镜像,graph 捕获期一次构建、replay 期复用,若设备侧序列长度变化而镜像未同步,稀疏注意力 kernel 会读到错误元数据;seq_lens_cpu_int 的 clamp(min=1) 兜底逻辑需要额外关注。
2) 版本绑定风险:统一到 torch.ops.npu.* 依赖 sgl-kernel-npu 对应算子版本(PR#689/#699/#708),且 extra_ops_loader.py 移除了 metadata op 的加载后校验,旧 .so 环境下错误会延迟到运行时才暴露。
3) 跨平台辐射:allocation.py 的条件扩展为 _is_hip or _is_npu,影响所有 NPU 平台的 last_loc 计算,不限于 DSV4,虽理论上更安全但覆盖范围大于本 PR 目标。
4) 配置读取改动:xpu_backend.py 改为 get_exec().deterministic,依赖运行时上下文已初始化,若在构造期不可用可能导致 XPU 启动回归。
5) 测试覆盖有限:单测仅补一个字段断言,缺少对 host metadata op 的定向单测;性能数据没有 PR 前基线,收益难以量化。

影响面集中在 NPU + DSV4/DSpark 路径:对部署 DeepSeek-V4(含 DSpark 投机解码)的 NPU 用户,D2H 同步消除可改善 overlapped 调度与吞吐,精度保持不变(aime26 0.9667)。统一命名空间降低了 normal/draft 双路径的 kernel 维护成本,但要求 sgl-kernel-npu 升级到对应算子版本,形成版本绑定。allocation.py 的改动会辐射所有 NPU triton 分配路径(行为更安全)。xpu_backend.py 的改动只影响 XPU 上的确定性推理开关读取方式。团队侧,NPU 后端与 sgl-kernel-npu 的协作模式更清晰,DSV4 多硬件适配脉络(XPU/NPU/AMD)进一步收敛。

核心路径变更(DSV4 稀疏注意力后端) 依赖 sgl-kernel-npu 新算子版本 缺少 PR 前性能基线 单测覆盖有限

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论