Kairos-50M 在昇腾 NPU 上的适配

Kairos-50M（mldi-lab/Kairos_50m）——一款自适应时间序列基础模型（arXiv:2509.25826），专为零样本预测设计，已通过 torch_npu 完成昇腾 NPU 适配。

模型描述

Kairos-50M 是一个拥有 5000 万参数的时间序列基础模型，在 PreSTS 语料库（超过 3000 亿个时间点）上完成预训练。该模型采用动态分块分词器（大小混合专家 MoE）、由 FFT 频谱特征驱动的实例级动态旋转位置编码（DRoPE），以及带有 eager attention 机制的 T5 风格编码器-解码器架构。模型可针对长达 2048 个点的输入上下文，输出 9 分位数（0.1 至 0.9）预测结果。

架构信息（来自 config.json）：

model_type：kairos，d_model：512，编码器 6 层，解码器 6 层（分为 2 个解码器段），num_heads：8，d_kv：64，d_ff：2048
动态分块：input_patch_size：128，input_patch_stride：128，levels：3，MoE n_activated_experts：3，n_null_experts：2
DRoPE：position_embedding_type：instance_wise_rope，instance_rope_input_feature_dim：128，rope_init：exp，scale_method：log
context_length：2048，prediction_length：每段 64（最大单次推理范围 2 × 64 = 128）
9 个分位数 [0.1, ..., 0.9]，约 5013 万参数，fp32 精度

NPU 适配说明

Kairos 官方模型代码（来自 github.com/foundation-model-research/Kairos 的 tsfm 包）已内置到 tsfm/ 目录中。该架构对 NPU 天然友好——T5 attention 采用 eager matmul 实现（不依赖 flash-attention 或 SDPA），动态分块、MoE 路由（topk/bincount）、InstanceNorm（nanmean）以及复数旋转位置编码（torch.polar/view_as_complex）均可在 NPU 上原生运行。

唯一一处源码修改（patches/kairos_ascend_npu.diff，位于 tsfm/model/kairos/modeling_kairos.py::fft_process）：NPU 不支持对 complex64 类型的 rfft 输出执行 torch.abs 操作（报错 aclnnAbs 错误码 161002），因此幅度改为手动计算为 sqrt(real^2 + imag^2)（与 torch.abs 相比最大误差 4.8e-6）。其余部分零代码改动。

vLLM-Ascend 0.18 未对自定义的 kairos 编码器-解码器架构提供注册支持，因此推理服务采用 torch_npu + FastAPI 方案部署。

环境

Ascend 910（64 GB 高带宽内存），CANN 8.5.1
Python 3.11.14，torch 2.9.0 / torch-npu 2.9.0.post1，transformers 4.57.6
einops 0.8.2，jaxtyping，fastapi / uvicorn

使用方法

1. 下载模型权重

python3 download_model.py     # -> /data/.cache/models/mldi-lab/Kairos_50m

权重文件从 GitCode HF 镜像 获取（https://ai.gitcode.com/hf_mirrors/mldi-lab/Kairos_50m，通过 resolve API 指向 LFS CDN，使用 12 条并行字节范围连接），并以 hf-mirror.com 作为备用源；sha256 校验值与 HF LFS oid（c68503d9…，200,551,636 字节）比对验证。

2. 启动服务

KAIROS_DEVICE=npu:1 PORT=8057 ./start_serve.sh
# or:
python3 server.py --port 8057

3. 推理

import json, urllib.request

payload = json.dumps({
    "past_values": [[...], [...], [...]],   # [C, T]  (each row = one series, T >= 2048)
    "prediction_length": 64,                 # horizon <= 128
}).encode()
req = urllib.request.Request("http://localhost:8057/v1/forecast",
                             data=payload, headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req) as resp:
    out = json.loads(resp.read())
print(out["prediction"])     # [C, 9, L]  quantile forecasts

4. 命令行界面

python3 make_test_data.py
python3 inference.py --input-json demo_series.json --device npu:1

基准测试（Ascend 910，npu:1）

上下文长度	预测步长	延迟（NPU p50）	CPU p50	显存占用
2048	64	~78 毫秒	~260 毫秒	~398 MB（进程）

首次请求包含 aclnn 算子编译（约 0.5 秒）；稳态阶段已实现内核缓存。NPU 与 CPU（fp32，相同检查点）预测结果的 余弦相似度 = 1.0000000000 / 最大绝对误差 ≤ 5.72e-6，约 3.5 倍加速 —— 详见 verify_npu_vs_cpu.py。

文件结构

inference.py — KairosEngine（库 + 命令行接口）
server.py — FastAPI /v1/forecast、/health 接口
download_model.py — 权重下载 + sha256 校验
make_test_data.py / test_server.py / verify_npu_vs_cpu.py — 演示与验证脚本
tsfm/ — 内置官方 Kairos 模型代码（含 1 处 NPU 补丁）
patches/kairos_ascend_npu.diff — 唯一的源码改动

权重文件不包含在本仓库中；请使用 download_model.py 进行下载。