TRELLIS.2-4B 昇腾 NPU 适配（子模型 smoke 推理）

1. 模型简介

TRELLIS.2（microsoft/TRELLIS.2-4B）是微软开源的 4B 参数 3D 资产生成模型，实现 image → latents → mesh / 3DGS 的端到端 pipeline，基于 O-Voxel 稀疏体素结构与 Flow-Matching Transformer，可生成最高 1536³ 分辨率的 PBR 3D 资产。

本仓库自包含以下三个子模块（本次 NPU 适配对象）：

适配目标：任意子模型可在昇腾 NPU（Ascend 910）上完成一次 forward（输入合成 latent tensor，输出 shape / slat 流），并对 NPU 与 CPU 的输出做数值与耗时对比。

NPU 适配要点

原始实现依赖 CUDA 专属库（flex_gemm 稀疏卷积、flash_attn 稀疏注意力、o_voxel / cumesh 网格后处理）。本次适配：

1. 稀疏卷积：新增纯 torch 后端 trellis2_src/trellis2/modules/sparse/conv/conv_none.py （子流形稀疏卷积，通过 SPARSE_CONV_BACKEND=none 启用）。
2. 稀疏注意力：full_attn.py / windowed_attn.py 增加 torch 后端 （varlen 切块 + F.scaled_dot_product_attention，通过 SPARSE_ATTN_BACKEND=torch 启用）。
3. CUDA 库去依赖：o_voxel / cumesh / flex_gemm 改为可选导入； o_voxel 不可用时 VAE 解码器直接返回原始解码流（不做 mesh 后处理）。
4. 权重按官方 *.json（模型结构）+ *.safetensors（参数）格式通过 trellis2.models.from_pretrained 加载。

2. 部署步骤

环境要求

安装依赖

pip install -r requirements.txt

源码与权重

# TRELLIS.2 源码（含 NPU 适配 patch，已随仓库提供：trellis2_src/）
# 权重（约 10-16G，本仓库不包含，需自行下载到 /data/dl/trellis2/）：
#   ckpts/shape_dec_next_dc_f16c32_fp16.{json,safetensors}
#   ckpts/shape_enc_next_dc_f16c32_fp16.{json,safetensors}
#   ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.{json,safetensors}

权重目录为 /data/dl/trellis2/ckpts（可在 inference.py 中通过 CKPT_DIR 修改）。

3. 推理示例

3.1 CPU 运行

python inference.py --device cpu            # 保存 cpu_output.txt + cpu_run.log + cpu_outputs.npz

3.2 NPU 运行

python inference.py --device npu            # 保存 npu_output.txt + npu_run.log + npu_outputs.npz

两个设备使用相同 seed（12345）与相同合成输入（256 个 latent 点，每子模块 forward 计时 3 次取平均），保证可对比性。

3.3 结果对比

python compare.py                           # 生成 assets/compare_result.txt（6 项指标）

运行输出示例（NPU）

=== TRELLIS.2-4B submodule smoke inference @ npu ===
[shape_dec] output shape: [80954, 7] | avg time: 825.41 ms | cum time: 825.41 ms | mem peak: 1054.9 MiB
[shape_enc] output shape: [63, 32] | avg time: 275.99 ms | cum time: 1101.40 ms | mem peak: 695.7 MiB
[slat_flow] output shape: [256, 32] | avg time: 110.86 ms | cum time: 1212.25 ms | mem peak: 2561.2 MiB
=== TOTAL cumulative avg time on npu: 1212.25 ms ===

4. 实测 6 指标对比表（NPU vs CPU，seed=12345，256 点）

分项耗时（3 次平均）

分项显存峰值

指标说明

• 数值类指标（1/2/3/5）按各子模块输出 flatten 后逐元素对齐计算（shape 不一致的子模块按较短长度截断）。
• 形状一致率（4）为任务特化指标：三个子模块中输出 shape 与 CPU 完全一致的占比。
• 耗时比（6）为 3 次 forward 平均耗时之比，<1 表示 NPU 更快。
• 偏差主要来自 fp16/bf16 低精度计算与稀疏卷积坐标去重顺序差异；用于 pipeline 前向 验证（数值合理、shape 正确、无算子 fallback 报错）已达成。