Toto-2.0-313m 在昇腾 NPU 上的部署

Datadog Toto-2.0-313m 是 Toto 2.0 家族中参数量为 313M 的成员，属于面向多变量预测的时间序列基础模型（技术报告：arXiv:2605.20119，TIME 基准：arXiv:2602.12147）。该模型为 u-μP 缩放、仅解码器、分块（patch） Transformer 架构，采用时间轴（因果）与变量轴（全连接）交替注意力机制，配备分位数输出头（9 个节点，覆盖 0.1 至 0.9，以 pinball 损失训练），支持连续块掩码（CPM） 以实现单次并行解码，内置稳健的 arcsinh 输入缩放器以及残差 MLP 分块投影。Datadog 将其定位为通用型强基线检查点——在 GIFT-Eval 上位列基础模型前三。

本次适配通过 torch_npu 将模型部署于 Ascend 910 NPU 上，用于基于服务的概率预测。由于 vLLM-Ascend 0.18 不支持此类自定义时间序列架构，故直接采用独立的 Toto2Model（来自 DataDog/toto，内置版本）在 torch_npu 上原样运行——模型源码零修改；仅剥除了内置包中的 GluonTS 封装层。

模型

属性	值
参数量	313M（fp32 safetensors 格式，1.17 GB）
架构	仅解码器 Transformer，`d_model=1024`，`num_layers=24`，`num_heads=16`，`qk_dim/v_dim=64`，`d_ff=2736`，`patch_size=32`，xPos RoPE，SwiGLU MLP，τ-规则残差（`residual_mult=0.75`），MuP（按维度缩放）
注意力	时间/变量交替分组（`layer_group_size=24`，末层为变量轴，`heads_per_group=1`）
输出	9 个分位数 `[0.1 … 0.9]`，形状为 `[9, batch, n_variate, horizon]`
任务	零样本多变量时间序列预测（可观测性 / GIFT-Eval / TIME）

环境

Ascend 910 ×2（单卡 64 GB HBM），CANN 8.5.1
Python 3.11，torch 2.9.0 + torch-npu 2.9.0.post1
unit-scaling==0.3.5（PyPI）+ 内置 dd_unit_scaling + 内置 toto2（剥离 GluonTS）
fastapi / uvicorn / pydantic

快速开始

# 1. download weights (ModelScope, sha256-verified)
python3 download_model.py

# 2. generate a demo test series
python3 make_test_data.py

# 3. run a one-shot forecast on NPU
python3 inference.py --device npu:1 --num-variates 3 --time-steps 1024 --prediction-length 96

# 4. start the HTTP service (default port 8040, device npu:1)
./start_serve.sh 8040 npu:1

服务 API

POST /v1/forecast

{
  "series": [[1.0, 2.0, ...], [...], ...],   // [variate, time_steps]; 1D accepted; NaN allowed
  "prediction_length": 96,                   // future steps (1..4096)
  "decode_block_size": 768,                  // CPM block size (multiple of 32)
  "has_missing_values": true
}

响应：{quantiles: {"0.1": […], …, "0.9": […]}, median, prediction_length, n_variates, latency_ms}，其中每个分位数都是一个 [n_variate, horizon] 数组。

测试用例与结果（Ascend 910, npu:1）

配置	延迟（HTTP p50）	备注
3 变量 / 上下文 1024 / 预测 96	约 120 ms	分位数单调
3 变量 / 上下文 512 / 预测 96	约 110 ms
1 变量 / 上下文 4096 / 预测 336	约 135 ms	长上下文
1 变量 / 上下文 1024 / 预测 1024	约 250 ms	KV-cache 块解码

模型加载：约 6 秒；进程 HBM 约 1.6 GB（npu:1）
NPU 与 CPU 对比（fp32），在 2 个变量 × 512 步 × 预测 96 下：每个分位数的 max_abs_diff ≤ 1e-4，cosine ≥ 0.9999999，中位数 corr ≈ 1.0 → 通过
稳态计算（2 变量 / 上下文 1024 / 预测 96）：得益于 24 层 transformer， NPU 相较于 CPU 在单次预测延迟上实现了显著加速。

NPU 适配说明

引擎：torch_npu。vLLM-Ascend 0.18（317 架构注册表）中没有针对该自定义时序模型的条目，因此直接加载独立的 Toto2Model。
零源码补丁：原生 F.scaled_dot_product_attention → aclnnFlashAttentionScore； fp64 因果补丁缩放器、asinh、分位数头以及 torch.distributions 均在 NPU 上原生运行。
内置依赖包（toto2/、dd_unit_scaling/）：toto2 的 Toto2GluonTSModel （GluonTS 包装器，会引入 lightning）已被剥离——参见 patches/toto2_ascend_npu.patch。 dd_unit_scaling（u-μP 层包装器）为内置包；上游的 unit-scaling 来自 PyPI。
NaN 处理：被掩码的 NaN 位置在输入时用零填充，以免污染因果缩放器的 cumsum 计算；模型仍会通过 target_mask 将它们视为缺失值。

交付物

inference.py（torch_npu 引擎 + 命令行界面）· server.py（FastAPI /v1/forecast）· download_model.py（ModelScope + sha256）· make_test_data.py · verify_npu_vs_cpu.py · test_server.py · start_serve.sh · toto2/ + dd_unit_scaling/（内置）· patches/toto2_ascend_npu.patch · requirements.txt · assets/（截图）

本仓库不包含模型权重——请通过 download_model.py 下载（来源：ModelScope Datadog/Toto-2.0-313m，sha256 5ee60664…aa2f9）。