[achena666]

这个 PR 做了什么

DiffSynth-Engine 目前只能跑在 CUDA 上：设备探测、通信后端、注意力实现、低精度量化四处都绑定了 CUDA 专用接口。本 PR 增加 Ascend NPU 后端，让 Qwen-Image 系列管线可以在昇腾上单卡与多卡运行，并支持 FB-Cache、序列并行、W8A8 量化三个特性。

设计原则：不改变任何现有默认行为。 所有新增分支都以 IS_NPU 为条件，CUDA 路径与改动前逐行一致；测试里的新增开关全部有默认值，不设置时等价于上游行为。

验证环境：openEuler 24.03 / CANN 9.2.0 / torch 2.8.0+cpu / torch_npu 2.8.0.post5 / transformers 4.53.0。

---

改动清单

1. utils/platform.py

• 新增 IS_NPU 探测（torch_npu 可导入且 torch.npu.is_available()），导入失败时安全降级
• empty_cache() 增加 NPU 分支

2. models/basic/attention.py（核心）

• npu_ulysses_attention()：all_to_all（聚合序列 / 切分 head）→ npu_fusion_attention → all_to_all 还原
• npu_sp_gather_attention()：ring_degree > 1 时的默认路径，把完整 K/V all_gather 到各 rank，Q 保持本地分片
• npu_ring_attention()：真正的环形 P2P 实现（NPU_RING_IMPL=ring 启用）
• long_context_attention() 在 NPU 上分发到上面三条路径

为什么必须改：AttnType.TORCH_EFFICIENT 指定 SDPA 的 mem-efficient 后端，torch_npu 未实现 → 自动回退 CPU → 输出张量留在 CPU，后续 HCCL 集合通信报 create_config ... error code 4 / DIST invalid type。这个错误信息把注意力故障伪装成通信故障，很难定位。

关于 yunchang 的 NPU ring 实现：ring_npu_flash_attn_forward 用 if step <= comm.rank 作为计算条件——这是 causal 分支，需要 not is_causal 守护——并且从不调用 update_out_and_lse 合并各分块的 softmax 统计量，实际只返回本地分片。实测 rank0 的 ring 输出与「只用本地 KV」的注意力完全相等（max_err=0），整条管线出的图与预期是两张。本 PR 的 npu_ring_attention 复用 yunchang 的 RingComm 与 update_out_and_lse，只重写循环，数值上与全量注意力一致（seq 64/512/2048，最大误差 0.004，即 bf16 精度）。

默认仍用 gather：环形每块先舍入到 bf16 再合并，40 步扩散会使端到端结果变差；且 2 卡下两者总通信量相同（都是 N-1 份 KV），环形的收益是计算与通信重叠。可用 NPU_RING_IMPL=ring 切换。

3. utils/parallel.py

• nccl → hccl（NPU 时）
• torch.cuda.set_device / synchronize 增加 NPU 分支
• device_id 只传给 NCCL（HCCL 不接受该参数）
• master_port 支持环境变量 DSNP_MASTER_PORT，避免端口冲突

4. utils/fp8_linear.py（W8A8 在线量化）

• 新增 npu_quant_linear()：激活走 npu_dynamic_quant(dst_type=292)（per-token，融合算子），
  权重离线做 per-channel float8_e4m3fn，矩阵乘走 npu_quant_matmul（fp8 × fp8 → bf16）
• NPU_QUANT_DTYPE=fp8|int8 切换位宽，默认 fp8
• NPU_QUANT_SKIP（默认 proj_out,to_out）按模块名保留 bf16 权重
• scale 存在 weight.npu_quant_scale（Parameter 属性）而非按 id(weight) 建表：
  fused LoRA 会改写 weight.data 而 Parameter 对象不变，按 id 索引会取到过期 scale
• 顺手修复：fp8_inference 的 yield 改为 try/finally，异常时也能还原 F.linear

FP8 在 Ascend 上的入口（值得记录，容易踩）：

接口	结论
torch.matmul(fp8, fp8)	不支持
npu_add_quant_matmul	算子支持 fp8+e8m0，但 mx 模式要求 transposeX1=true，torch 层固定传 false，eager 路径不可用
npu_transpose_quant_batchmatmul	3D perm 语义在本版本有 batch 轴推断问题
npu_quant_matmul + npu_dynamic_quant(dst_type=292)	可用

dst_type 是 torch_npu 的内部编号，等于 aclDataType + 256：292 = ACL_FLOAT8_E4M3FN(36)，296 = ACL_FLOAT4_E2M1(40)，290 = HIFLOAT8。

scale 的形状/类型有硬性要求：per-channel 权重 scale 必须 1D float32（2D 会被要求转 int64），per-token 激活 scale 必须 float32。

5. models/qwen_image/qwen_image_dit_fbcache.py

• forward 补齐缺失参数（edit / context_latents / zero_cond_t / entity_*，txt_seq_lens 改名 text_seq_lens）——原实现停留在旧 API，开启即 TypeError
• 新增 warmup_steps（前 N 步强制计算）与 max_skip_streak（限制连续跳过），默认 0 与上游行为一致
• 缓存判断包在 sequence_parallel 上下文内
• 跳过决策在 SP 组内做 all_reduce(MAX) 同步：各卡独立判断会导致「有的卡跳过、有的卡继续」，下一步集合通信时死锁

6. pipelines/qwen_image.py

• 构造 QwenImageDiTFBCache 时补传 use_zero_cond_t（原先只有非缓存分支传了，2511 走缓存会算错）

7. tests/common/test_case.py

• get_input_image() 用 with ... as image: return image.convert("RGB") 强制加载像素。多设备管线会把输入图 pickle 后传给 worker，PIL 惰性加载会在此时读已关闭的文件句柄并报 broken PNG file
• assertImageEqual() 无论通过与否都打印 SSIM，便于 CI 日志引用（判定与阈值不变）

8. tests/test_pipelines/test_qwen_image_edit_plus.py

• 权重路径、并行度、设备、量化开关改为环境变量，全部有默认值，不设置时与上游行为一致
• 并行模式下如果 device 带设备索引（"npu:0"）直接报错：mp.spawn 的每个 worker 都会先 set_device(rank)，带索引会让所有 worker 挤到 0 号卡
• 移除 width=1328, height=1328：参考图是 1248×832（由输入图比例自动推算），传入固定尺寸会与参考图尺寸不符，structural_similarity 直接抛异常

---

验证

功能正确性

项目	结果
SP 注意力 vs 全量参考（seq 64/512/2048）	最大绝对误差 0.004（bf16 精度），余弦 > 0.99999
融合注意力 vs SDPA	余弦 0.99999994，最大绝对误差 0
单卡确定性（同配置重跑）	SSIM = 1.000000
CFG 并行 vs 单卡	SSIM = 1.0000
量化单层余弦	fp8 0.9993 / int8 0.9999

端到端（Qwen-Image-Edit-2509，1248×832，40 步）

配置	SSIM	耗时
单卡 bf16	0.9458	136 s
2 卡 Ulysses	0.9507	105 s
2 卡 ring 组（全量 KV）	0.9507	94 s
2 卡 + W8A8 FP8	0.9481	87 s
2 卡 + W8A8 INT8（对照）	0.7914	84 s

上表数字来自回退旧参考图参数的内部跑法（DSNP_BASELINE_ERA 一类开关），未包含在本 PR 中。
按本 PR 的默认配置直接跑，该用例是 0.659——与 main 相同，见下节。

量化的取舍：FP8 相对 bf16 掉 0.0026，速度 1.2×，权重显存减半；INT8 掉 0.159——per-token 量程被激活离群值拉爆后小值全被压成 0，而 FP8 的指数位覆盖得住。这也是默认选 fp8 的原因。

---

已知问题（不在本 PR 范围内）

test_qwen_image_edit_plus.py 在 main 上即为 0.659，本 PR 未改变这一行为。

tests/data/expect/qwen_image/qwen_image_edit_plus.png 生成于 2025-09-30，此后上游有三处改动会改变数值：

• 4215ce4：条件图 resize 由 LANCZOS 改为 PIL 默认（BICUBIC），并给调度器加了 shift_terminal=0.02
• 0590444：VAE 分块解码相关
• 另有 GELU 实现调整

把这三处临时回退后，单卡可以回到 0.9458——但这属于「复现旧参考图」的工作，不适合塞进「支持 NPU」的 PR。建议单独开 issue 讨论是否更新参考图。

width=1328/height=1328 与参考图尺寸（1248×832）不符也源于此，本 PR 只移除了这两个参数，没有改动断言。

---

不包含的内容

• 没有修改任何 CUDA 路径的行为
• 没有引入新的依赖（torch_npu 是可选导入，缺失时 IS_NPU 为 False）
• 没有提交为复现旧参考图而用的临时开关（DSNP_BASELINE_ERA、DSNP_SSIM_THRESHOLD 等）

[achena666]

Superseded by #265 — my changes have been merged into that branch:
fbcache fixes, W8A8 FP8 quantization, and the torch_npu attention fallback.

Verified on 2 x Ascend950PR (1248x832, 40 steps, reference-image parameters
restored): 0.96017 SSIM PASS. Also reverted _modulate to the original
expression, since torch.addcmul costs ~0.014 SSIM in bf16.

Branch: https://github.com/achena666/DiffSynth-Engine/tree/npu-merge-265
Closing to keep a single Ascend NPU PR.