VideoFlexTok d18-d28 on Ascend NPU (torch_npu 2.9.0.post1)

1. 简介

本文档记录 EPFL-VILAB/videoflextok_d18_d28（VideoFlexTok，灵活长度的由粗到细视频 tokenizer）在华为昇腾 Ascend910 NPU 上的适配与真机验证结果。

模型是一个视频离散 tokenizer / 自编码器：RGB 视频片段 [C, T, H, W] 先经 VidTok 因果 3D-VAE 压成潜变量，再由 18 层 FlexTransformer 编码器编成 FSQ 离散 token（levels=[8,8,8,5,5,5]，codebook_size=64000），解码端用一个 28 层、宽度 dim=1792 的 rectified-flow（流匹配）扩散解码器 从 token 重建视频。d18_d28 = 编码器 18 层 / 解码器 28 层（比同族 d18_d18 解码器更深更宽），工作分辨率 256×256（VAE 潜变量 5×32×32），每片段 17 帧。权重 model.safetensors 约 10.04 GB（10,779,872,388 字节，832 个张量），fp32 推理。

适配要点：

• 不需要 trust_remote_code：官方 videoflextok 包（git+https://github.com/apple/ml-videoflextok）+ 仓库 config.json（hydra 装配）即可加载；VidTok VAE 权重已含在同一个 model.safetensors 里，无需额外下载外部基模型。
• 核心昇腾改造：上游注意力用 torch.nn.attention.flex_attention + create_block_mask 并经 torch.compile 走 CUDA/Triton inductor 后端，昇腾无该后端。本 inference.py 在导入后将该路径 monkeypatch 成数学等价的稠密掩码 scaled_dot_product_attention（把 block-mask 物化成布尔 [b,1,N,N] 张量），从而让同一份权重在 npu:0 上跑通，无任何 CUDA kernel。
• VidTok VAE 用 scaled_dot_product_attention 与因果 3D 卷积，昇腾原生支持；强制 fp32、HF32 关闭以保证精度。
• d28 相对 d18_d18 的差异（本次校准点）：解码器 depth 18→28、dim 1152→1792，分辨率 128→256，VAE 潜变量 [5,16,16]→[5,32,32]；编码器与 FSQ 码本不变。ckpt 严格加载 832 张量、missing=0 / unexpected=0，解码器实测 28 个 block，与 config 完全一致。
• 本模型特有验证点：token 形状 [1, 5, 256]、token 分布合理、预训练重建 PSNR 决定性优于同架构随机初始化。

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/EPFL-VILAB/videoflextok_d18_d28
• 权重下载地址（GitCode 镜像）：https://ai.gitcode.com/hf_mirrors/EPFL-VILAB/videoflextok_d18_d28
• 上游推理代码：https://github.com/apple/ml-videoflextok

2. 验证环境

• NPU：Ascend910，2 逻辑卡（npu-smi 25.5.5，65536 MB HBM/卡），本次用 card 0
• 推理引擎：torch_npu
• 权重路径：~/.cache/models/videoflextok_d18_d28/model.safetensors（约 10.04 GB，832 个张量）
• 任务类型：video tokenizer / video autoencoder（离散 token 编码 + 流匹配解码重建）

3. 环境准备与权重获取

export PATH=/usr/local/python3.11.14/bin:$PATH
# 上游包（务必 --no-deps，避免其 pyproject 的 torch==2.8.0 覆盖 pod 的 torch_npu 2.9.0）
pip install --user --no-deps --no-build-isolation "git+https://github.com/apple/ml-videoflextok"
pip install --user --no-deps diffusers==0.35.1
pip install --user hydra-core omegaconf mup einops
# 前置检查
python -c "import torch, torch_npu; assert torch.npu.is_available(); print('npu OK')"

权重获取（inference.py 已自包含此逻辑，评委空缓存可直接重跑）：走 HF_ENDPOINT=https://hf-mirror.com 的 snapshot_download，Xet CDN 域名已加入 no_proxy 直连。model.safetensors（10,779,872,388 字节，sha256:d54e409a…）文件较大；若 hf-mirror 拥塞，可用 GitCode 镜像的 git-lfs batch API 取签名 CDN 直链、多路 Range 并行 curl 拉取，合并后 sha256 校验一致即可。

4. 推理运行

已验证通过的命令（套卡锁 + 固定物理卡 + 编译缓存进 /tmp）：

flock -w 7200 /tmp/npu_card0.lock bash -c '
  export PATH=/usr/local/python3.11.14/bin:$PATH
  ASCEND_RT_VISIBLE_DEVICES=0 ASCEND_CACHE_PATH=/tmp/ascend_cache_videoflextok_d18_d28 \
  ASCEND_PROCESS_LOG_PATH=/tmp/ascend_log ASCEND_WORK_PATH=/tmp/ascend_work \
  python inference.py --device npu'

inference.py 会：加载权重（Gate1 严格加载）→ 校验 config 落地（Gate2）→ 用合成结构化视频做编码/解码 smoke → 跑 Gate3 三证 → 采集性能，并把结果写 /tmp/videoflextok_d18_d28_npu_results.json。

5. Smoke 验证

上述命令在昇腾 NPU 上的真实输出（原样粘贴）：

[npu-adapt] flex_attention -> dense-mask SDPA patch applied
======================================================================
Loading EPFL-VILAB/videoflextok_d18_d28 on npu:0 (fp32)
[Gate1] strict load (832 tensors): missing=0 (real/non-buffer 0) unexpected=0
[Gate2] chunk_size=OK; vae_video_sizes=OK; video_preprocess_args.size=OK; encoder.depth=OK; decoder.depth=OK; fsq.levels=OK
        FSQ codebook_size=64000 (=prod(levels))

[demo] structured synthetic video shape=(3, 17, 256, 256) (C,T,H,W) range[-1.000,1.000]

======================================================================
SMOKE (pretrained autoencode):
  token seqs: 1  shape=[1, 5, 256]  dtype=torch.int64
  token id range: [9, 63953] (valid codebook 0..63999)
  recon shape=(1, 3, 17, 256, 256) range[-1.000,1.000]
  reconstruction: MSE=0.01295  PSNR=24.898 dB

======================================================================
GATE 3 (a) reconstruction — pretrained vs random-init (same architecture):
  PRETRAINED : PSNR=24.898 dB  MSE=0.01295  tokrange[9,63953]
  RANDOM-INIT: PSNR=11.085 dB  MSE=0.31156  tokrange[32036,32036]
  => PSNR gap = +13.813 dB  (pretrained DECISIVELY wins)

[Gate3-b] stability x2: tokens bit-exact=False, agreement=99.61% (1275/1280); encoder+FSQ deterministic (fixed VAE latents)=True; pretrained PSNR reproducible=24.90dB
[Gate3-c] token dist: 1215 unique ids / 1280 total, all in [0,63999]=True; recon in [-1,1]=True

======================================================================
PERF: full autoencode (tokenize + 20-step flow decode), fp32, npu:0
  avg=9938.2ms  min=9920.7  max=9959.8  p50=9934.8  p90=9955.5  p95=9956.7
  throughput=0.101 clips/s  peak_HBM=28072.9 MB

验证结果：

• 权重严格加载 832 张量，missing=0 / unexpected=0（Gate1 通过）；解码器实测 28 个 block，与 d28 config 一致。
• config 关键项（chunk_size、vae_video_sizes=[5,32,32]、编码 depth=18/解码 depth=28、video_preprocess_args.size=256、fsq.levels）全部落到实例（Gate2 通过）。
• 17 帧 256×256 视频编码成 token 序列 [1, 5, 256]（int64），id 全落在合法码本 [0, 63999]。
• 流匹配解码重建输出 [1, 3, 17, 256, 256]，取值在 [-1, 1]，重建 PSNR=24.90 dB。
• 二次运行完全一致：Gate1（832/0/0）、重建 PSNR=24.898 dB、gap +13.813 dB 逐次复现。

6. 性能参考

测试条件：单片段 [3,17,256,256]，一次「编码 + 20 步流匹配解码」为一次调用；warmup=5，正式 20 次，每次前后 torch.npu.synchronize()；fp32，npu:0。

耗时主要在 20 步扩散解码（每步一次完整 28 层 Transformer 前向 + CFG 双分支）；相比同族 d18_d18（解码 18 层、128×128，约 6.7 s/clip、10 GB HBM），d28 因解码器更深更宽、分辨率翻倍，单次约 9.9 s、峰值 HBM 约 28 GB。稳态 p50/p90 波动 <0.5%，无明显 CANN 尖刺。

7. 精度评测

本模型是无类别 GT 的视频 tokenizer/autoencoder，按 Gate-3「三证」评测（无权威在线基线，比对方式为预训练 vs 同架构随机初始化的重建对照 + 确定性 + token 分布）：

8. 适配截图

9. 注意事项

最容易踩的坑：上游注意力的 flex_attention + torch.compile 在昇腾上没有后端，不改就无法前向。

实际失败特征如下：

• 现象：模型能加载，但一进入 tokenize/detokenize 的 Transformer 前向即报错，或 torch.compile(flex_attention) 触发 inductor/Triton 相关失败。
• 关键报错：torch.nn.attention.flex_attention / create_block_mask 依赖 CUDA-Triton inductor 后端，昇腾无对应实现。
• 位置：videoflextok/model/layers/attention.py::FlexAttention.forward（self.flex_attention = torch.compile(flex_attention)）与 videoflextok/model/preprocessors/flex_seq_packing.py::create_block_mask_cached。

原因不是「算子精度问题」或「dtype 不对」，而是 flex_attention 这一整条 CUDA/Triton 专属 的稀疏注意力路径在昇腾上缺后端。

当前环境的可用处理方式（inference.py 已内置，导入后、加载权重前生效）：

# 1) 把 block-mask 物化成稠密布尔掩码
def _patched_create_block_mask(mask_fn, B, H, M, N, device="cpu", _compile=False):
    q = torch.arange(M, device=device).view(M,1).expand(M,N)
    k = torch.arange(N, device=device).view(1,N).expand(M,N)
    mask = mask_fn(0,0,q,k).view(1,1,M,N) if B is None else \
           torch.stack([mask_fn(b,0,q,k) for b in range(B)],0).unsqueeze(1)
    return _DenseMask(mask)
fsp.create_block_mask_cached = _patched_create_block_mask   # 必须在实例化前 patch

# 2) FlexAttention.forward 改走稠密掩码 SDPA（等价、无 torch.compile）
attn_mod.FlexAttention.forward = _flex_forward  # F.scaled_dot_product_attention(q,k,v, attn_mask=mask, scale=self.scale)

其余注意事项：

1. 上游包必须 --no-deps --no-build-isolation 安装：其 pyproject.toml 硬钉 torch==2.8.0，若带依赖安装会覆盖 pod 的 torch_npu 2.9.0 栈；--no-build-isolation 还能避开构建期联网拉 setuptools 而卡住。hydra-core 1.3.x 需搭配 antlr4-python3-runtime==4.9.3，否则 Could not deserialize ATN。
2. diffusers 不在预装列表：videoflextok.flow_matching 依赖 FlowMatchEulerDiscreteScheduler，需 pip install --no-deps diffusers==0.35.1。
3. d28 是 ~10 GB fp32 ckpt：Gate1 会先建一份 CPU 副本做严格加载，inference.py 在进入推理前已 del 该副本 + state_dict 并 gc.collect()，再实例化推理模型搬到 NPU，避免同时驻留多份大模型。
4. VidTok 因果 3D 卷积在 NPU 上非确定：同一输入两次 tokenize 潜变量有 ~1e-2 抖动，导致约 0.4% 的 token 在 FSQ 量化边界翻转；这是卷积规约顺序问题，torch.use_deterministic_algorithms(True) / HCCL_DETERMINISTIC 均无法完全消除，但不影响重建质量与 token 分布，属可接受现象。
5. ASCEND_CACHE_PATH 目录须在 import 时预建，否则会伪装成 HF32/ACL error 500001（inference.py 顶部已 makedirs）。
6. token 序列形状为 [1, 5, 256]（本 ckpt n_max=256 寄存器）。

10. 标签

#NPU #Ascend