vllm-ascend 逐请求 NaN 检测内核 `_num_nans_kernel`：Triton 实现、Ascend 适配与精度验证-平芜编程栈

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

【免费下载链接】vllm-ascend

Community maintained hardware plugin for vLLM on Huawei Ascend

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

导读

本文围绕 vllm-ascend 在采样链路上用于统计 logits 中 NaN 数量的 Triton 内核_num_nans_kernel展开，介绍其逐请求（逐行）计数语义、分块遍历算法、在 Ascend 上替换上游 CUDA 实现的 patch 机制，以及覆盖低精度输入与特殊值的精度测试。读完本文，你将掌握该内核的调用方式、参数约束、与上游 vLLM 实现的差异，以及如何在昇腾 NPU 上验证其逐位精确的计数行为。

功能定位：为采样链路提供 logits 质量信号

在 vLLM 推理中，logits 中出现 NaN 通常意味着数值不稳定（如溢出、除零或权重异常），需要被采样链路感知。_num_nans_kernel的职责是逐请求（逐行）统计 logits 中的 NaN 数量，输出 INT32 计数张量，供采样链路上报num_nans_in_logits指标。该指标既可用于观测（经nans_to_dict组织为{req_id: count}字典，配合raise_if_nan_logits可在VLLM_RAISE_ON_LOGIT_NANS=1时直接抛错），也可作为recompute_scheduler等调度组件判断 logits 质量并触发重算的依据。

在生产环境中，该内核通过 vllm_ascend/patch/worker/patch_v2/patch_triton.py 将上游vllm.v1.worker.gpu.metrics.logits.get_num_nans整体替换为本实现（包装函数为get_num_nans），源码位于 vllm_ascend/ops/triton/v2/metrics/num_nans.py。

内核设计：逐行并行 + 分块掩码遍历

计数公式

对每个请求行r，内核计算该行 logits 中满足isnan的元素个数：

$$num_nans[r] = \sum_{i=0}^{vocab_size-1} \mathbb{1}\left[isnan(logits[r, i])\right]$$

算法流程

每个 program 只负责一个请求行，tl.program_id(0)即请求行索引，grid 尺寸为(num_reqs,)。单行内的处理流程如下：

分块遍历：按常量BLOCK_SIZE分块遍历该行 vocab 维度。vocab_size不是BLOCK_SIZE整数倍时，循环多执行一轮以覆盖尾部剩余元素。
掩码加载：以mask = block < vocab_size掩码加载行内元素，越界位置以other=0填充——0 不是 NaN，不会污染计数结果。
数据类型上转：加载数据统一转为 FLOAT32，调用 CANN libdevice 的isnan判定 NaN，结果转为 INT32。
块内归约：tl.sum对块内 INT32 结果求和，并累加到标量num_nans。
写回：循环结束后，将该行总计数写入num_nans_ptr[r]。

从源码看，num_nans.py 的实现与文档描述完全一致：标量累加使用num_nans += tl.sum(is_nan).to(tl.int32)，因此输出为精确整数计数，无浮点累加误差。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
logits_ptr	输入	二维`[num_reqs, vocab_size]`logits 张量。	FLOAT32	ND
logits_stride	属性	`logits`的 stride(0)，即行跨度。	INT	-
num_nans_ptr	输出	一维`[num_reqs]`，每行的 NaN 计数。	INT32	ND
vocab_size	属性	vocab 维度长度。	INT	-
BLOCK_SIZE	属性	常量，分块遍历宽度，生产固定为 8192。	constexpr INT	-

包装函数get_num_nans（num_nans.py）对调用方隐藏了大部分细节：它断言logits为二维张量，读取num_reqs与vocab_size，固定使用BLOCK_SIZE = 8192，分配num_reqs长度的 INT32 输出张量后直接以(num_reqs,)的 grid 启动内核，并传入logits.stride(0)作为行跨度。

与上游实现的差异与 patch 机制

差异根源

上游 vLLM 的vllm/v1/worker/gpu/metrics/logits.py中，libdevice 从torch._inductor.runtime.triton_helpers导入，会解析到 CUDA 符号，在 Ascend 上导致内核编译失败；且其启动路径依赖triton.experimental.gluon.nvidia，而 triton-ascend 不提供该模块。

vllm-ascend 版本则直接导入triton.language.extra.cann.libdevice，内核逻辑与上游一致，仅 libdevice 来源不同。

替换方式

vllm_ascend/patch/worker/patch_v2/patch_triton.py 通过整体替换完成接入：

sampler.get_num_nans = get_num_nans rejection_sampler.get_num_nans = get_num_nans

即同时替换sampler与rejection_sampler两个模块中的get_num_nans。patch 注释明确说明：该补丁属于 Triton-Ascend 后端兼容性适配，待编译器与 Triton Ascend toolkit 支持上游实现后可以回退（详见 patch/init.py 中对vllm.v1.worker.gpu.metrics.logits.libdevice的说明）。

调用示例

以下代码完整复现文档示例，同时给出生产包装入口与直接启动内核两种方式（需在昇腾 NPU 环境运行）：

import torch from vllm_ascend.ops.triton.v2.metrics.num_nans import ( _num_nans_kernel, get_num_nans, ) device = "npu" num_reqs, vocab_size = 4, 4096 logits = torch.randn(num_reqs, vocab_size, dtype=torch.float32, device=device) logits[0, :10] = float("nan") # 请求 0 注入 10 个 NaN # 方式一：生产包装函数（patch 后 sampler 实际调用的入口） num_nans = get_num_nans(logits) # 方式二：直接启动内核（与精度 UT 一致） num_nans = torch.empty(num_reqs, dtype=torch.int32, device=device) _num_nans_kernel(num_reqs,), num_nans, vocab_size, BLOCK_SIZE=8192, ) torch.npu.synchronize() print(num_nans.cpu()) # 期望 [10, 0, 0, 0]

约束说明

logits必须为二维 FLOAT32；内核按logits_stride + 偏移访问行内元素，因此支持行跨度大于vocab_size的排布（如带 padding 的 batch 张量）。
num_nans必须为一维 INT32，长度等于num_reqs。
BLOCK_SIZE生产固定为 8192：vocab_size小于BLOCK_SIZE时靠掩码保护，大于时多轮迭代，两种情况均正确。
每个 program 只处理一个请求行，grid 尺寸应为(num_reqs,)。
输出为精确整数计数，无浮点累加误差，对比时应逐位相等（rtol=0, atol=0）。

精度验证：120 组参数组合与特殊值专项

精度 UT 位于 tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_num_nans.py，直接测试本 Ascend 内核，与 CPU 参考实现（torch.isnan(logits).sum(dim=-1).to(torch.int32)）逐位比对（rtol=0, atol=0）。

参数组合覆盖

数据类型：dtype ∈ {fp32, bf16, fp16}。内核每次加载后执行.to(tl.float32)上转再判 NaN，测试验证低精度输入经上转后 NaN 不丢失、不误产生。
NaN 布局：行头连续（prefix）与逐行随机散布（scatter，每行使用独立随机位置集），防止位置性偏差掩盖误判。
形状：(num_reqs, vocab_size)覆盖 128/5000/8192/10000/16384，对应尾部掩码、非 2 幂、单块/多块 + 部分尾部掩码组合。
NaN 比例：frac_nan ∈ {0.0, 0.1, 0.5, 1.0}，其中0.0与1.0覆盖全无 NaN 与全 NaN 两种边界。

以上组合共 120 组。数据在 CPU 上以 fp32 构造并注入 NaN 后再转目标 dtype，参考值按转换后的张量计算，从而精确检验低精度转换过程对 NaN 位型的影响。

特殊值专项

test_num_nans_special_values_not_counted将 ±Inf/±0.0/max-finite/min-normal/subnormal 等特殊位型平铺满行，要求计数为 0；另一行在其上叠加散布的 NaN，要求恰好计数。该测试专门防止 CANNisnan退化为“无穷/指数类”判断——即避免把 ±Inf、-0.0、subnormal 等误判为 NaN，确保判定严格遵循 IEEE NaN 语义。

运行命令：

pytest -sv tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_num_nans.py

总结

_num_nans_kernel是 vllm-ascend 采样链路中一个小而关键的算子：以逐行并行、分块掩码遍历的方式在昇腾 NPU 上精确统计 logits 中的 NaN 数量，并通过 patch 机制无缝替换上游 CUDA 取向的get_num_nans。其实现充分体现了硬件插件适配的典型模式——保持上游内核逻辑与语义不变，仅替换 libdevice 来源与启动路径，同时以逐位等价的精度 UT 守住数值正确性底线。对于需要在采样链路上观测或处置 logits 数值异常的开发者，该内核及其测试用例是一份可直接参照的实现范本。