[FEATURE] Run Wan2.2 pipelines on Ascend NPU

[yx0716]

Description

Enable the Wan2.2 pipelines to run on Ascend NPU, single-card and multi-card, with a minimal correctness-focused change set. All fixes reuse the existing platform architecture (current_platform, capability gates, queue_with_cpu); CUDA behavior is unchanged by construction — every new branch is gated on platform or tensor device type.

Motivation

TeleFuser already ships a platform abstraction (telefuser/platforms/, CustomOp.forward_npu, dist_backend="hccl"), and the wan22 pipeline code itself is platform-clean. In practice, a handful of CUDA defaults outside that abstraction still prevented NPU execution: an unconditional channels_last_3d activation cast, a "cuda" device-mesh default, device-IPC-dependent worker queues, a hard import triton, and "cuda"-allocated P2P buffers. This PR removes those blockers.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Code refactoring
• Documentation update
• Other (please describe):

Changes Made

• models/wan_video_vae.py, distributed/vae_spatial.py: gate activation-side channels_last_3d casts to CUDA tensors (weight-side conversion was already cuDNN-gated; NPU rejects that memory format with ERR01007).
• distributed/device_mesh.py: create_device_mesh_from_config defaults device_type to current_platform.device_type (12 of 14 call sites rely on the default; the "cuda" default crashed torch DeviceMesh on NPU hosts with integer modulo by zero). Explicit arguments keep working.
• worker/parallel_worker.py: enable the existing queue_with_cpu marshalling by default on non-CUDA platforms and mirror it on the result path (torch_npu cross-process tensor sharing raises devptr INTERNAL ASSERT FAILED when rebuilding queued NPU tensors). Worker dispatch already moves inputs to the local device, so the contract is unchanged; CUDA keeps direct device-queue transport.
• ops/fp8_attention.py → kernel/triton/fp8_attention.py: move the @triton.jit fused FP8 QKV kernels into kernel/triton/ and import them lazily inside quantize_fp8_qkv after its existing CUDA-only validation, matching the established ops -> kernel.triton dispatch pattern. The wan_video / minimax DiT import chains no longer require triton to be installed.
• distributed/pp_comm.py, tests/unit/distributed/test_pp_comm.py: allocate P2P receive buffers on current_platform.device_type instead of hard-coded "cuda", and parameterize the test tensor device so the suite runs on CUDA, NPU, and CPU hosts.
• examples/wan_video/wan22_t2v_5b.py: default the pipeline device to the detected platform so the example runs unmodified on CUDA and NPU hosts (resolves to "cuda" on CUDA as before).

Testing

• Unit tests pass (pytest tests/)
• Manual testing performed
• Benchmarks added/updated (if applicable)

Test commands:

# CPU host (triton not installed): import chain and unit tests
python -c "from telefuser.pipelines.wan_video.wan22_ti2v import Wan22TI2VPipeline"
pytest tests/unit/distributed/test_pp_comm.py tests/unit/worker -q
pytest tests/unit -q

# Ascend 910B (CANN 8.2, torch 2.9.0 + torch_npu 2.9.0.post1, triton not installed)
python examples/wan_video/wan22_t2v_5b.py --gpu_num 1   # unmodified example; device auto-detected
python examples/wan_video/wan22_t2v_5b.py --gpu_num 4   # cfg_degree=2 x sp_ulysses_degree=2 over hccl

Results on Ascend 910B: single-card Wan2.2-TI2V-5B text-to-video succeeds end to end (bf16, TORCH_SDPA, unipc; semantically correct output); the MODEL_CPU_OFFLOAD variant is bit-identical to the no-offload run (same seed). 4-card denoising completes with workers exiting cleanly, using the existing eager ulysses fallbacks. tests/unit/distributed/test_pp_comm.py passes 12/12 (previously 6 device-related failures on non-CUDA hosts), and tests/unit/models/test_wan_video_sol_attention.py passes with CUDA-only cases skipped. CUDA hosts execute identical code paths as before (device-mesh default resolves to "cuda", queue_with_cpu stays opt-in, activation casts and buffer devices are unchanged on CUDA).

Checklist

• Code follows the project's coding standards (ruff)
• Pre-commit hooks pass (pre-commit run --all-files — ruff + ruff-format; also covered by the lint CI job)
• All tests pass (pytest tests/)
• New tests added for new functionality (platform-parameterized test_pp_comm.py)
• Documentation updated (docstrings; no README/AGENTS impact — a platform guide is planned as follow-up)
• Commit messages are clear and descriptive
• PR title follows the convention: [TYPE] Brief description

Related Issues

N/A

Additional Notes

Performance parity is intentionally out of scope: with everything on eager fallbacks, multi-card currently shows limited speedup over single-card on NPU (per-layer eager all_to_all over hccl and the serial VAE decode dominate). A follow-up PR will address NPU-optimized paths behind the existing dispatch points. offload/async_offload.py remains CUDA-only (auto-disabled elsewhere); Wan2.2 A14B (MoE dual-DiT) shares these code paths but has not been exercised on NPU hardware yet.

GPU Architecture Support

• SM80 (Ampere, Ada Lovelace)
• SM90 (Hopper H100)
• SM100+ (Blackwell)

No CUDA kernels are added or modified — the fused FP8 QKV kernels are moved verbatim and their SM90 targeting is unchanged — so no architecture box is checked.

Device support added by this PR: Ascend NPU (verified on Atlas 910B, CANN 8.2, torch 2.9.0 + torch_npu 2.9.0.post1, via the existing telefuser/platforms/ NPU platform). CUDA devices are unaffected.

Performance Impact

None on CUDA (all new branches are platform- or device-gated; CUDA executes the previous code paths).

Ascend 910B baseline established by this PR — Atlas 910B (8× 910B2 64 GB), CANN 8.2, torch 2.9.0 + torch_npu 2.9.0.post1. All numbers are produced by the unmodified example examples/wan_video/wan22_t2v_5b.py (bf16, TORCH_SDPA, unipc, 50 denoising steps, MODEL_CPU_OFFLOAD, cfg_scale=5.0 — i.e. two DiT forwards per denoising step; under cfg parallelism the two guidance branches run on separate ranks). Times are end-to-end generate (one-time init excluded); every row is reproducible with the exact command shown:

Workload	Command	1 card	2 cards (cfg=2)	4 cards (cfg=2 × ulysses=2)
480p, 121 frames	python examples/wan_video/wan22_t2v_5b.py --gpu_num N	180.9 s	120.8 s (1.50×)	87.0 s (2.08×)
720p, 121 frames	python examples/wan_video/wan22_t2v_5b.py --gpu_num N --resolution 720p	477.8 s	—	173.9 s (2.75×)

Measured stage split at 480p / 121 frames (from run-log timestamps):

Stage	1 card	2 cards	4 cards
Text encode (incl. weight onload)	≈10.6 s	≈9 s	≈8.8 s
Denoise, 50 steps	136 s (2.73 s/step)	67 s (2.03×)	40 s (3.4×)
DiT weight staging + VAE decode + result return	≈34 s	≈45 s	≈38 s