From cf3248c63fa93f3d5753b9eb06add4e7947c5889 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 00:45:14 +0800 Subject: [PATCH 01/10] fix(models): gate channels_last_3d activations to CUDA devices The VAE weight-side conversion to channels_last_3d is already guarded by cuDNN availability, but the activation-side casts in CausalConv3d.forward and the spatial-parallel halo conv were unconditional. NPU and CPU reject channels_last_3d activations, so Wan2.2 VAE forward failed on Ascend with ERR01007 ("NPU contiguous operator only supported contiguous memory format"). Route non-CUDA tensors through standard contiguous instead. Verified: Wan2.2-TI2V-5B text-to-video smoke on Ascend 910B2 (CANN 8.2, torch_npu 2.9.0) completes end to end; CUDA path unchanged. --- telefuser/distributed/vae_spatial.py | 6 +++++- telefuser/models/wan_video_vae.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/telefuser/distributed/vae_spatial.py b/telefuser/distributed/vae_spatial.py index 85f2f410..ebeb30fe 100644 --- a/telefuser/distributed/vae_spatial.py +++ b/telefuser/distributed/vae_spatial.py @@ -128,7 +128,11 @@ def _spatial_causal_conv3d_forward( if any(padding): tensor = F.pad(tensor, padding) tensor = _exchange_height_halo(module, tensor, module._height_halo_size) - tensor = tensor.contiguous(memory_format=torch.channels_last_3d) + if tensor.device.type == "cuda": + tensor = tensor.contiguous(memory_format=torch.channels_last_3d) + else: + # channels_last_3d activations are only supported by cuDNN; NPU/CPU require standard contiguous. + tensor = tensor.contiguous() return F.conv3d( tensor, module.weight, diff --git a/telefuser/models/wan_video_vae.py b/telefuser/models/wan_video_vae.py index 22b4bc7c..941f8056 100644 --- a/telefuser/models/wan_video_vae.py +++ b/telefuser/models/wan_video_vae.py @@ -119,7 +119,11 @@ def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) - x = x.contiguous(memory_format=torch.channels_last_3d) + if x.device.type == "cuda": + x = x.contiguous(memory_format=torch.channels_last_3d) + else: + # channels_last_3d activations are only supported by cuDNN; NPU/CPU require standard contiguous. + x = x.contiguous() return super().forward(x) From e9fd11636d80022837fe4b4797e9d5bf5579c0d7 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 00:45:14 +0800 Subject: [PATCH 02/10] fix(distributed): default device mesh to the detected platform create_device_mesh_from_config defaulted device_type to "cuda" and 12 of 14 call sites rely on the default, so parallel denoising on NPU crashed inside torch DeviceMesh with "integer modulo by zero" (zero visible CUDA devices). Resolve the default from current_platform; explicit arguments keep working. Verified: 4-card (cfg=2 x ulysses=2) Wan2.2-TI2V-5B run on Ascend 910B2 builds the [cfg, ulysses] mesh over hccl; on CUDA the default resolves to "cuda" as before. --- telefuser/distributed/device_mesh.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/telefuser/distributed/device_mesh.py b/telefuser/distributed/device_mesh.py index 8281327e..ce3e69ff 100644 --- a/telefuser/distributed/device_mesh.py +++ b/telefuser/distributed/device_mesh.py @@ -13,10 +13,11 @@ from torch.distributed.device_mesh import DeviceMesh from telefuser.core.config import ParallelConfig +from telefuser.platforms import current_platform from telefuser.utils.logging import logger -def create_device_mesh_from_config(parallel_config: ParallelConfig, device_type: str = "cuda") -> DeviceMesh: +def create_device_mesh_from_config(parallel_config: ParallelConfig, device_type: str | None = None) -> DeviceMesh: """Create PyTorch DeviceMesh from ParallelConfig. Mesh dimensions are built in order: DP -> CFG -> SP (ring, ulysses) -> PP -> TP @@ -24,11 +25,13 @@ def create_device_mesh_from_config(parallel_config: ParallelConfig, device_type: Args: parallel_config: Parallel configuration with degrees for each dimension - device_type: Device type ("cuda" or "cpu") + device_type: Device type ("cuda", "npu", or "cpu"); defaults to the current platform's device type Returns: PyTorch DeviceMesh instance with named dimensions """ + if device_type is None: + device_type = current_platform.device_type _validate_parallel_config(parallel_config) sp_degree = parallel_config.sp_ulysses_degree * parallel_config.sp_ring_degree From 7d5c8f457b3e2f3aff758ddbe76cc9181c600078 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 00:45:14 +0800 Subject: [PATCH 03/10] fix(worker): marshal parallel-worker queue tensors through CPU off CUDA Parallel workers exchange request and result tensors over torch.multiprocessing queues, which relies on device IPC. torch_npu cross-process sharing raised "devptr INTERNAL ASSERT FAILED ... entry in cache has missing shared_ptr" when rebuilding queued NPU tensors. Enable the existing queue_with_cpu marshalling by default on non-CUDA platforms, mirror it on the result path, and move results back to the stage device in the main process. Worker dispatch already moves inputs to the local device, so the contract is unchanged and CUDA keeps direct device-queue transport. Verified: 4-card Wan2.2-TI2V-5B denoising on Ascend 910B2 completes with workers exiting cleanly. --- telefuser/worker/parallel_worker.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 2c160544..9adb8450 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -158,6 +158,9 @@ def _worker_loop( y = tensor_output_channel.send(y) # Always output results when world_size=1 if world_size == 1 or rank == 0: + if current_platform.device_type != "cuda": + # Queue transport without reliable device IPC: marshal results through CPU. + y = to_device(y, "cpu") queue_out.put(y) except Exception as e: import traceback @@ -206,7 +209,8 @@ def __init__( self.device_ids = list(range(self.world_size)) self.name: str = f"Parallel Worker {stage.name}" - self.queue_with_cpu: bool = parallel_config.queue_with_cpu + # Queue transport requires CPU marshalling on platforms without reliable device IPC (e.g. NPU). + self.queue_with_cpu: bool = parallel_config.queue_with_cpu or current_platform.device_type != "cuda" self.timeout: int = parallel_config.timeout self._lifecycle_lock = threading.Lock() self._failed = False @@ -319,6 +323,9 @@ def _wait_result(self, method_name: str) -> Any: reason = f"{method_name} failed: {result}" self._mark_failed(reason) raise RuntimeError(f"ParallelWorker:{self.name} {reason}") from result + if current_platform.device_type != "cuda": + # CPU-marshalled queue results move back to the stage device for downstream consumers. + result = to_device(result, self._stage.device) return result def enable_metrics(self, registry: Any | None = None) -> None: From 74dd3abd613804509849884ecc243a9d61205abf Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 00:45:22 +0800 Subject: [PATCH 04/10] refactor(ops): move fused FP8 QKV Triton kernels into kernel.triton ops/fp8_attention.py defined @triton.jit kernels at module scope, making "import triton" unconditional for every consumer of the wan_video and minimax DiT import chains even though the fused path is CUDA-only by contract. Move the two kernels and their launcher to telefuser/kernel/triton/fp8_attention.py and import them lazily inside quantize_fp8_qkv after its existing CUDA validation, matching the established ops -> kernel.triton dispatch pattern (see ops/rotary.py and kernel/__init__.py). The pure-torch quantize/dequantize helpers are unchanged. Verified: wan22 pipeline import and Wan2.2-TI2V-5B smoke succeed on an NPU host without triton installed; tests/unit/models/test_wan_video_sol_attention.py passes with CUDA-only cases skipped. --- telefuser/kernel/triton/fp8_attention.py | 119 +++++++++++++++++++++++ telefuser/ops/fp8_attention.py | 110 +-------------------- 2 files changed, 123 insertions(+), 106 deletions(-) create mode 100644 telefuser/kernel/triton/fp8_attention.py diff --git a/telefuser/kernel/triton/fp8_attention.py b/telefuser/kernel/triton/fp8_attention.py new file mode 100644 index 00000000..a620f374 --- /dev/null +++ b/telefuser/kernel/triton/fp8_attention.py @@ -0,0 +1,119 @@ +"""Fused block-scaled FP8 Q/K/V quantization Triton kernels.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _quantize_qkv_fp8_stage1( + q, + k, + v, + q_out, + k_out, + q_scale, + k_scale, + v_scale, + tokens: tl.constexpr, + heads: tl.constexpr, + head_dim: tl.constexpr, + block: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_head = tl.program_id(1) + batch = batch_head // heads + head = batch_head % heads + token_offsets = block_idx * block + tl.arange(0, block) + dim_offsets = tl.arange(0, head_dim) + valid = token_offsets < tokens + offsets = ((batch * tokens + token_offsets[:, None]) * heads + head) * head_dim + dim_offsets[None, :] + q_values = tl.load(q + offsets, mask=valid[:, None], other=0.0).to(tl.float32) + k_values = tl.load(k + offsets, mask=valid[:, None], other=0.0).to(tl.float32) + v_values = tl.load(v + offsets, mask=valid[:, None], other=0.0).to(tl.float32) + + q_s = tl.maximum(tl.max(tl.max(tl.abs(q_values), axis=1), axis=0), 1.0e-6) / 448.0 + k_s = tl.maximum(tl.max(tl.max(tl.abs(k_values), axis=1), axis=0), 1.0e-6) / 448.0 + scale_offset = (batch * tl.cdiv(tokens, block) + block_idx) * heads + head + tl.store(q_scale + scale_offset, q_s) + tl.store(k_scale + scale_offset, k_s) + tl.store(q_out + offsets, q_values / q_s, mask=valid[:, None]) + tl.store(k_out + offsets, k_values / k_s, mask=valid[:, None]) + + v_s = tl.max(tl.abs(v_values), axis=0) / 448.0 + v_scale_offsets = (batch * heads + head) * head_dim + dim_offsets + tl.atomic_max(v_scale + v_scale_offsets, v_s) + + +@triton.jit +def _quantize_qkv_fp8_stage2_v( + v, + v_out, + v_scale, + tokens: tl.constexpr, + heads: tl.constexpr, + head_dim: tl.constexpr, + block: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_head = tl.program_id(1) + batch = batch_head // heads + head = batch_head % heads + token_offsets = block_idx * block + tl.arange(0, block) + dim_offsets = tl.arange(0, head_dim) + valid = token_offsets < tokens + input_offsets = ((batch * tokens + token_offsets[:, None]) * heads + head) * head_dim + dim_offsets[None, :] + output_offsets = ((batch * heads + head) * head_dim + dim_offsets[None, :]) * tokens + token_offsets[:, None] + scale_offsets = (batch * heads + head) * head_dim + dim_offsets + scale = tl.maximum(tl.load(v_scale + scale_offsets), 1.0e-6 / 448.0) + values = tl.load(v + input_offsets, mask=valid[:, None], other=0.0).to(tl.float32) + tl.store(v_out + output_offsets, values / scale[None, :], mask=valid[:, None]) + + +def quantize_fp8_qkv_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Launch the fused Q/K/V quantization kernels on validated CUDA inputs.""" + batch, tokens, heads, head_dim = q.shape + blocks = triton.cdiv(tokens, block_size) + q_out = torch.empty(q.shape, device=q.device, dtype=torch.float8_e4m3fn) + k_out = torch.empty_like(q_out) + v_storage = torch.empty((batch, heads, head_dim, tokens), device=q.device, dtype=torch.float8_e4m3fn) + q_scale = torch.empty((batch, blocks, heads), device=q.device, dtype=torch.float32) + k_scale = torch.ones_like(q_scale) + v_scale = torch.zeros((batch, heads, head_dim), device=q.device, dtype=torch.float32) + grid = (blocks, batch * heads) + _quantize_qkv_fp8_stage1[grid]( + q, + k, + v, + q_out, + k_out, + q_scale, + k_scale, + v_scale, + tokens, + heads, + head_dim, + block_size, + num_warps=8, + num_stages=1, + ) + _quantize_qkv_fp8_stage2_v[grid]( + v, + v_storage, + v_scale, + tokens, + heads, + head_dim, + block_size, + num_warps=8, + num_stages=1, + ) + v_out = v_storage.permute(0, 3, 1, 2) + return q_out, k_out, v_out, q_scale, k_scale, v_scale diff --git a/telefuser/ops/fp8_attention.py b/telefuser/ops/fp8_attention.py index 7a2f36d5..eeec0f20 100644 --- a/telefuser/ops/fp8_attention.py +++ b/telefuser/ops/fp8_attention.py @@ -4,77 +4,10 @@ import torch import torch.nn.functional as F -import triton -import triton.language as tl FP8_ATTENTION_BLOCK_SIZE = 64 -@triton.jit -def _quantize_qkv_fp8_stage1( - q, - k, - v, - q_out, - k_out, - q_scale, - k_scale, - v_scale, - tokens: tl.constexpr, - heads: tl.constexpr, - head_dim: tl.constexpr, - block: tl.constexpr, -): - block_idx = tl.program_id(0) - batch_head = tl.program_id(1) - batch = batch_head // heads - head = batch_head % heads - token_offsets = block_idx * block + tl.arange(0, block) - dim_offsets = tl.arange(0, head_dim) - valid = token_offsets < tokens - offsets = ((batch * tokens + token_offsets[:, None]) * heads + head) * head_dim + dim_offsets[None, :] - q_values = tl.load(q + offsets, mask=valid[:, None], other=0.0).to(tl.float32) - k_values = tl.load(k + offsets, mask=valid[:, None], other=0.0).to(tl.float32) - v_values = tl.load(v + offsets, mask=valid[:, None], other=0.0).to(tl.float32) - - q_s = tl.maximum(tl.max(tl.max(tl.abs(q_values), axis=1), axis=0), 1.0e-6) / 448.0 - k_s = tl.maximum(tl.max(tl.max(tl.abs(k_values), axis=1), axis=0), 1.0e-6) / 448.0 - scale_offset = (batch * tl.cdiv(tokens, block) + block_idx) * heads + head - tl.store(q_scale + scale_offset, q_s) - tl.store(k_scale + scale_offset, k_s) - tl.store(q_out + offsets, q_values / q_s, mask=valid[:, None]) - tl.store(k_out + offsets, k_values / k_s, mask=valid[:, None]) - - v_s = tl.max(tl.abs(v_values), axis=0) / 448.0 - v_scale_offsets = (batch * heads + head) * head_dim + dim_offsets - tl.atomic_max(v_scale + v_scale_offsets, v_s) - - -@triton.jit -def _quantize_qkv_fp8_stage2_v( - v, - v_out, - v_scale, - tokens: tl.constexpr, - heads: tl.constexpr, - head_dim: tl.constexpr, - block: tl.constexpr, -): - block_idx = tl.program_id(0) - batch_head = tl.program_id(1) - batch = batch_head // heads - head = batch_head % heads - token_offsets = block_idx * block + tl.arange(0, block) - dim_offsets = tl.arange(0, head_dim) - valid = token_offsets < tokens - input_offsets = ((batch * tokens + token_offsets[:, None]) * heads + head) * head_dim + dim_offsets[None, :] - output_offsets = ((batch * heads + head) * head_dim + dim_offsets[None, :]) * tokens + token_offsets[:, None] - scale_offsets = (batch * heads + head) * head_dim + dim_offsets - scale = tl.maximum(tl.load(v_scale + scale_offsets), 1.0e-6 / 448.0) - values = tl.load(v + input_offsets, mask=valid[:, None], other=0.0).to(tl.float32) - tl.store(v_out + output_offsets, values / scale[None, :], mask=valid[:, None]) - - def quantize_fp8_qkv( q: torch.Tensor, k: torch.Tensor, @@ -86,46 +19,11 @@ def quantize_fp8_qkv( raise ValueError("q, k, and v must share shape [B, T, H, D]") if not (q.is_cuda and q.is_contiguous() and k.is_contiguous() and v.is_contiguous()): raise ValueError("fused FP8 QKV quantization requires contiguous CUDA tensors") - batch, tokens, heads, head_dim = q.shape - if head_dim != 128: + if q.shape[-1] != 128: raise ValueError("fused FP8 QKV quantization requires head dimension 128") - blocks = triton.cdiv(tokens, FP8_ATTENTION_BLOCK_SIZE) - q_out = torch.empty(q.shape, device=q.device, dtype=torch.float8_e4m3fn) - k_out = torch.empty_like(q_out) - v_storage = torch.empty((batch, heads, head_dim, tokens), device=q.device, dtype=torch.float8_e4m3fn) - q_scale = torch.empty((batch, blocks, heads), device=q.device, dtype=torch.float32) - k_scale = torch.ones_like(q_scale) - v_scale = torch.zeros((batch, heads, head_dim), device=q.device, dtype=torch.float32) - grid = (blocks, batch * heads) - _quantize_qkv_fp8_stage1[grid]( - q, - k, - v, - q_out, - k_out, - q_scale, - k_scale, - v_scale, - tokens, - heads, - head_dim, - FP8_ATTENTION_BLOCK_SIZE, - num_warps=8, - num_stages=1, - ) - _quantize_qkv_fp8_stage2_v[grid]( - v, - v_storage, - v_scale, - tokens, - heads, - head_dim, - FP8_ATTENTION_BLOCK_SIZE, - num_warps=8, - num_stages=1, - ) - v_out = v_storage.permute(0, 3, 1, 2) - return q_out, k_out, v_out, q_scale, k_scale, v_scale + from telefuser.kernel.triton.fp8_attention import quantize_fp8_qkv_triton + + return quantize_fp8_qkv_triton(q, k, v, FP8_ATTENTION_BLOCK_SIZE) def quantize_fp8_per_block( From b6dea52ab1338f9078cb95d42f83b115e0df841d Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 00:45:22 +0800 Subject: [PATCH 05/10] fix(distributed): allocate pipeline P2P buffers on the platform device recv, recv_latent, and recv_latent_async allocated receive buffers with device="cuda", which fails on NPU-only hosts before irecv can run. Allocate on current_platform.device_type (identical behavior on CUDA) and parameterize the test tensor device the same way so the suite runs on CUDA, NPU, and CPU hosts. Verified: tests/unit/distributed/test_pp_comm.py 12 passed on Ascend 910B2 (previously 6 device-related failures). --- telefuser/distributed/pp_comm.py | 7 ++++--- tests/unit/distributed/test_pp_comm.py | 18 +++++++++++------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/telefuser/distributed/pp_comm.py b/telefuser/distributed/pp_comm.py index 9c3b9085..5a4a2de7 100644 --- a/telefuser/distributed/pp_comm.py +++ b/telefuser/distributed/pp_comm.py @@ -24,6 +24,7 @@ import torch import torch.distributed as dist +from telefuser.platforms import current_platform from telefuser.utils.logging import logger @@ -107,7 +108,7 @@ def recv( if buffer is None: if shape is None: raise ValueError("Either buffer or shape must be provided") - buffer = torch.empty(shape, dtype=torch.float16, device="cuda") + buffer = torch.empty(shape, dtype=torch.float16, device=current_platform.device_type) buffer = buffer.contiguous() if async_op: @@ -266,7 +267,7 @@ def recv_latent(self, shape: tuple | None = None, dtype: torch.dtype = torch.bfl if shape is None: raise ValueError("recv_latent: shape must be provided") - buffer = torch.empty(shape, dtype=dtype, device="cuda") + buffer = torch.empty(shape, dtype=dtype, device=current_platform.device_type) buffer = buffer.contiguous() work = dist.irecv(buffer, self.recv_src, group=self._process_group) work.wait() @@ -300,7 +301,7 @@ def recv_latent_async(self, shape: tuple, dtype: torch.dtype = torch.bfloat16) - if self.is_first_stage: raise RuntimeError("recv_latent_async: First stage has no previous stage to receive from") - buffer = torch.empty(shape, dtype=dtype, device="cuda") + buffer = torch.empty(shape, dtype=dtype, device=current_platform.device_type) buffer = buffer.contiguous() work = dist.irecv(buffer, self.recv_src, group=self._process_group) return buffer, work diff --git a/tests/unit/distributed/test_pp_comm.py b/tests/unit/distributed/test_pp_comm.py index f247d001..99b52a7d 100644 --- a/tests/unit/distributed/test_pp_comm.py +++ b/tests/unit/distributed/test_pp_comm.py @@ -7,10 +7,14 @@ import torch.distributed as dist from telefuser.distributed.pp_comm import PipelineP2PComm +from telefuser.platforms import current_platform # Skip if distributed not available HAS_DISTRIBUTED = dist.is_available() +# Buffers follow the detected platform so the suite runs on CUDA, NPU, and CPU hosts alike. +DEVICE = current_platform.device_type + pytestmark = [ pytest.mark.skipif(not HAS_DISTRIBUTED, reason="Distributed not available"), pytest.mark.distributed, @@ -56,7 +60,7 @@ def test_send_on_last_stage_logs_warning(self): """Test that send on last stage logs warning and returns None.""" comm = PipelineP2PComm(None) # Single GPU, is_last_stage=True - tensor = torch.randn(1, 10, 512, device="cuda") + tensor = torch.randn(1, 10, 512, device=DEVICE) result = comm.send(tensor) assert result is None @@ -72,7 +76,7 @@ def test_send_recv_single_gpu(self): """Test send_recv on single GPU (no-op).""" comm = PipelineP2PComm(None) # Single GPU - send_tensor = torch.randn(1, 10, 512, device="cuda") + send_tensor = torch.randn(1, 10, 512, device=DEVICE) result = comm.send_recv(send_tensor) # On single GPU, recv_buffer is None since there's no previous stage @@ -114,7 +118,7 @@ def test_queue_send_on_last_stage(self): """Test queue_send on last stage does nothing.""" comm = PipelineP2PComm(None) # is_last_stage=True - tensor = torch.randn(1, 10, 512, device="cuda") + tensor = torch.randn(1, 10, 512, device=DEVICE) comm.queue_send(tensor) assert len(comm._ops) == 0 @@ -123,7 +127,7 @@ def test_queue_recv_on_first_stage(self): """Test queue_recv on first stage does nothing.""" comm = PipelineP2PComm(None) # is_first_stage=True - buffer = torch.randn(1, 10, 512, device="cuda") + buffer = torch.randn(1, 10, 512, device=DEVICE) comm.queue_recv(buffer) assert len(comm._ops) == 0 @@ -152,7 +156,7 @@ def test_send_latent_on_last_stage(self): """Test send_latent on last stage returns early.""" comm = PipelineP2PComm(None) # is_last_stage=True - tensor = torch.randn(1, 10, 512, device="cuda") + tensor = torch.randn(1, 10, 512, device=DEVICE) # Should not raise, just return comm.send_latent(tensor) @@ -182,7 +186,7 @@ def test_send_latent_async_on_last_stage(self): """Test send_latent_async on last stage returns None.""" comm = PipelineP2PComm(None) # is_last_stage=True - tensor = torch.randn(1, 10, 512, device="cuda") + tensor = torch.randn(1, 10, 512, device=DEVICE) result = comm.send_latent_async(tensor) assert result is None @@ -248,7 +252,7 @@ def test_actual_p2p_communication(self): if comm.is_first_stage: # Send tensor to next stage - send_tensor = torch.ones(1, 10, 512, device="cuda") * comm.rank + send_tensor = torch.ones(1, 10, 512, device=DEVICE) * comm.rank comm.send_latent(send_tensor) elif comm.is_last_stage: # Receive tensor from previous stage From 00a2124a9bc59b640d3fb0988006ae3030e00104 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 09:11:53 +0800 Subject: [PATCH 06/10] fix(worker): keep marshalled queue results device-agnostic The non-CUDA result path introduced for CPU-marshalled queues moved results back to self._stage.device inside _wait_result. That breaks worker unit tests on non-CUDA hosts (mocked stages resolve .device through the method proxy or to a MagicMock), and it is unnecessary: stage entry points already place their inputs, and the Wan VAE decode paths move latents themselves. Drop the move so _wait_result returns results untouched; workers still marshal results through CPU when device IPC is unavailable. Verified: tests/unit/worker passes on an Ascend host (same non-CUDA branch as CPU CI) and the 4-card Wan2.2-TI2V-5B smoke still completes; ruff clean. --- telefuser/worker/parallel_worker.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 9adb8450..01727299 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -323,9 +323,6 @@ def _wait_result(self, method_name: str) -> Any: reason = f"{method_name} failed: {result}" self._mark_failed(reason) raise RuntimeError(f"ParallelWorker:{self.name} {reason}") from result - if current_platform.device_type != "cuda": - # CPU-marshalled queue results move back to the stage device for downstream consumers. - result = to_device(result, self._stage.device) return result def enable_metrics(self, registry: Any | None = None) -> None: From 1b9642584ce689aab39bcb3d26f98c3ad126bc44 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 11:30:50 +0800 Subject: [PATCH 07/10] feat(examples): auto-detect device in the wan22 5B T2V example The example hardcoded device="cuda", so running it on an NPU host required editing the script. Default the pipeline device to current_platform.device_type so the same command runs unmodified on CUDA and NPU; on CUDA hosts this resolves to "cuda" as before. Verified: python examples/wan_video/wan22_t2v_5b.py --gpu_num 1 and --gpu_num 4 run unmodified on Ascend 910B (50 steps, 121 frames, default prompt). --- examples/wan_video/wan22_t2v_5b.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/wan_video/wan22_t2v_5b.py b/examples/wan_video/wan22_t2v_5b.py index 76cc1a56..573f1eb5 100644 --- a/examples/wan_video/wan22_t2v_5b.py +++ b/examples/wan_video/wan22_t2v_5b.py @@ -17,6 +17,7 @@ Wan22TI2VPipeline, Wan22TI2VPipelineConfig, ) +from telefuser.platforms import current_platform from telefuser.utils.utils import get_example_name from telefuser.utils.video import get_target_video_size_from_ratio, save_video @@ -80,7 +81,7 @@ def get_pipeline(parallelism: int = 1, model_root: str = PPL_CONFIG["model_root" ) # Create pipeline - pipe = Wan22TI2VPipeline(device="cuda", torch_dtype=torch.bfloat16) + pipe = Wan22TI2VPipeline(device=current_platform.device_type, torch_dtype=torch.bfloat16) # Configure pipeline pipe_config = Wan22TI2VPipelineConfig() From edf966dbbcceb18482f8f0a4d2a5201d498f03ad Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 02:09:18 +0800 Subject: [PATCH 08/10] fix(models): accept batched tensors in wan22 VAE parallel decode Wan22VideoVAE.decode's parallel branch called torch.stack on its input, assuming a list of per-video tensors, but VAEStage.decode_video passes a batched [B, C, T, H, W] tensor; the serial branch only worked because iterating a tensor yields its batch slices. Any enable_vae_parallel run of the Wan2.2 48-channel VAE therefore failed with TypeError regardless of platform. Normalize tensor inputs before stacking. Found while enabling spatially parallel VAE decode on Ascend 910B2; the parallel path now proceeds to communicator setup (further multi-communicator progress on that host is limited by its CANN driver, see branch notes). --- telefuser/models/wan22_video_vae.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/telefuser/models/wan22_video_vae.py b/telefuser/models/wan22_video_vae.py index 8ca32a35..6e70f21f 100644 --- a/telefuser/models/wan22_video_vae.py +++ b/telefuser/models/wan22_video_vae.py @@ -1403,7 +1403,11 @@ def decode( if self.parallelism > 1 and dist.is_initialized(): # tiled=True → tile_dist, tiled=False → 2d_split method = "tile_dist" if tiled else "2d_split" - hidden_states_tensor = torch.stack(hidden_states) + # The stage passes a batched [B, C, T, H, W] tensor; lists of per-video tensors are stacked. + if isinstance(hidden_states, torch.Tensor): + hidden_states_tensor = hidden_states + else: + hidden_states_tensor = torch.stack(hidden_states) return self.decode_parallel(hidden_states_tensor, device, method=method) # Single GPU processing From d34c842713880a6af8b8cb7a5a8614664ed633cf Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Tue, 1 Sep 2026 02:09:18 +0800 Subject: [PATCH 09/10] fix(worker): isolate HCCL socket ranges per parallel worker group Concurrent worker groups sharing the same NPU devices (e.g. denoising plus VAE workers) collide on HCCL's default data-plane socket range and fail comm init with EJ0003 ("IP address and port have been bound already"). Allocate a distinct HCCL_IF_BASE_PORT per spawned group on NPU platforms, mirroring the existing per-group MASTER_PORT allocation; other platforms are untouched. Verified: single-group 4-card Wan2.2-TI2V-5B regression on Ascend 910B2 passes with the env applied (33.4s generate, parity with baseline). --- telefuser/worker/parallel_worker.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index 01727299..b025d5ae 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -7,6 +7,7 @@ from __future__ import annotations import gc +import itertools import os import signal import threading @@ -35,6 +36,10 @@ _DISCARD_TENSOR_REFS = "__telefuser_discard_tensor_refs__" +# Each concurrent worker group needs a distinct HCCL socket range on NPU hosts; +# overlapping ranges fail comm init with EJ0003 (port already bound). +_hccl_group_counter = itertools.count() + def to_device(data: Any, device: str | torch.device) -> Any: """Recursively move data to target device.""" @@ -65,6 +70,7 @@ def _worker_loop( master_port: int, tensor_output_channel: WorkerTensorChannel | None = None, tensor_input_channels: tuple[WorkerTensorChannel, ...] = (), + hccl_if_base_port: int | None = None, ) -> None: """Worker process main loop. @@ -89,6 +95,11 @@ def _worker_loop( os.environ["WORLD_SIZE"] = str(world_size) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) + if hccl_if_base_port is not None: + # Give each worker group a distinct host-side HCCL socket base; + # overlapping ranges fail comm init with EJ0003 when several + # groups share the same devices (e.g. denoise + VAE workers). + os.environ["HCCL_IF_BASE_PORT"] = str(hccl_if_base_port) device_ids = parallel_config.device_ids device_id = rank if device_ids is not None: @@ -257,6 +268,7 @@ def __init__( master_port, self.tensor_output_channel, self.tensor_input_channels, + 60000 + 1024 * next(_hccl_group_counter) if current_platform.device_type == "npu" else None, ), nprocs=self.world_size, join=False, From 54a2b943fc07e80b0d372782bd9791dbcd15edb5 Mon Sep 17 00:00:00 2001 From: jinyx5 Date: Fri, 11 Sep 2026 16:02:51 +0800 Subject: [PATCH 10/10] docs: register Ascend NPU in the platform, installation, and wan_video guides Mirror the ROCm registration pattern for the NPU support validated by this branch, updating the guides introduced on main: - docs/en|zh/platforms.md: expand the "NPU and CPU" section with the native-fallback attention path (TORCH_SDPA, no triton/tf-kernel needed), forward_npu dispatch, HCCL with CPU-marshalled worker queues and per-group HCCL_IF_BASE_PORT isolation, and the validated wan22_t2v_5b.py entry point on Atlas 910B (CANN 8.2); note that Wan2.2 A14B is not yet exercised. - docs/en|zh/installation.md: add an Ascend NPU requirements row and a verification note (torch_npu selection, ASCEND_RT_VISIBLE_DEVICES, torch.cuda.is_available() False is expected). - examples/wan_video/README.md: add the Ascend requirements bullet and platform auto-detection / validated-hardware bullets to wan22_t2v_5b.py. Verification: python scripts/docs/prepare_cookbook.py build passes with rendered pages, assets, and anchors validated; the tests/docs unittest failures on this host are macOS /var tempdir symlink artifacts unrelated to these files. English additions stay within the 120-column style. --- docs/en/installation.md | 6 ++++++ docs/en/platforms.md | 23 +++++++++++++++++++---- docs/zh/installation.md | 5 +++++ docs/zh/platforms.md | 19 ++++++++++++++++--- examples/wan_video/README.md | 6 ++++++ 5 files changed, 52 insertions(+), 7 deletions(-) diff --git a/docs/en/installation.md b/docs/en/installation.md index a13c2e10..7bcefbed 100644 --- a/docs/en/installation.md +++ b/docs/en/installation.md @@ -12,6 +12,7 @@ distribution are installed separately. | PyTorch | 2.6 or newer | | CUDA toolkit | 12.8 or newer for the maintained CUDA development path | | ROCm | 7.x with a PyTorch `+rocm` build for AMD GPUs; see the ROCm note under verification | +| Ascend NPU | CANN 8.2 with matching `torch` and `torch_npu` builds; see the NPU note under verification | | GPU | Depends on the selected model; check its Cookbook guide | An example may impose stricter versions or GPU architecture requirements. In particular, locally built `tf-kernel` @@ -62,6 +63,11 @@ ROCm before CUDA. Examples ending in `_rocm.py` (for example `examples/wan_video are the validated entry points; see [Hardware Platforms](platforms.md) for per-platform capabilities and backend availability. +On Huawei Ascend hosts, install a `torch_npu` build matching your PyTorch version instead of the CUDA toolkit +path. The platform layer selects NPU when `torch_npu` imports and an Ascend device is visible +(`ASCEND_RT_VISIBLE_DEVICES` controls visibility); `torch.cuda.is_available()` printing `False` is expected there. +`examples/wan_video/wan22_t2v_5b.py` is the validated NPU entry point. + ## Model Checkpoints TeleFuser does not bundle model weights. The [Supported Models](supported_models.md) page links to each Cookbook diff --git a/docs/en/platforms.md b/docs/en/platforms.md index 50de62b7..353950e6 100644 --- a/docs/en/platforms.md +++ b/docs/en/platforms.md @@ -65,7 +65,22 @@ setup path. ## NPU and CPU -The NPU platform targets Huawei Ascend devices through `torch_npu` with the HCCL distributed backend. It is wired -into the platform and ops dispatch layers, but the maintained examples are validated on CUDA and, for select -examples, ROCm — validate on your target NPU before production use. The CPU platform is the fallback when no -accelerator is detected; it is intended for tests and for pipelines that explicitly request CPU execution. +The NPU platform targets Huawei Ascend devices through `torch_npu` with the HCCL distributed backend. + +- Attention uses `TORCH_SDPA` through the native fallback paths; no `tf-kernel`, `flash_attn`, `sageattention`, or + `triton` installation is required. +- The ops layer selects `forward_npu` where an op defines one and otherwise falls back to native PyTorch. No + NPU-optimized kernels are integrated yet, so pipelines run entirely on the native paths. +- Multi-card inference uses HCCL through the `hccl` backend string. Parallel-worker queues marshal tensors through + CPU because `torch_npu` has no reliable cross-process device IPC, and each spawned worker group receives a + distinct `HCCL_IF_BASE_PORT` so concurrent groups do not collide on HCCL's data-plane socket range. +- `torch.compile` is not validated on NPU; NPU examples run eager. +- The validated entry point is + [Wan2.2 TI2V-5B text-to-video](https://github.com/Tele-AI/TeleFuser/tree/main/examples/wan_video) + (`wan22_t2v_5b.py`), which auto-detects the platform and runs unmodified on an Atlas 910B (CANN 8.2, torch 2.9 + with a matching `torch_npu`): single-card, and four-card CFG × Ulysses parallelism over HCCL. Wan2.2 A14B shares + these code paths but has not been exercised on NPU hardware; validate other examples on your target NPU before + production use. + +The CPU platform is the fallback when no accelerator is detected; it is intended for tests and for pipelines that +explicitly request CPU execution. diff --git a/docs/zh/installation.md b/docs/zh/installation.md index a4df9c7e..2ec15661 100644 --- a/docs/zh/installation.md +++ b/docs/zh/installation.md @@ -11,6 +11,7 @@ | PyTorch | 2.6 或更高版本 | | CUDA Toolkit | 当前 CUDA 开发路径要求 12.8 或更高版本 | | ROCm | AMD GPU 使用 ROCm 7.x 与 PyTorch `+rocm` 构建，详见验证安装一节的说明 | +| 昇腾 NPU | CANN 8.2 搭配版本匹配的 `torch` 与 `torch_npu` 构建，详见验证安装一节的说明 | | GPU | 取决于所选模型，以对应 Cookbook 为准 | 具体示例可能要求更严格的软件版本或 GPU 架构。特别是本地构建的 `tf-kernel` 产物与其记录的 PyTorch、 @@ -61,6 +62,10 @@ AMD ROCm 主机应安装 PyTorch `+rocm` 构建，而非 CUDA Toolkit 路径。H `examples/wan_video/wan21_1_3b_text_to_video_rocm.py`）是已验证的入口；各平台能力与后端可用性见 [硬件平台](platforms.md)。 +华为昇腾主机应安装与 PyTorch 版本匹配的 `torch_npu`，而非 CUDA Toolkit 路径。当 `torch_npu` 可导入且存在 +可见昇腾设备（由 `ASCEND_RT_VISIBLE_DEVICES` 控制）时，平台层会选择 NPU；此时 `torch.cuda.is_available()` +输出 `False` 属预期行为。`examples/wan_video/wan22_t2v_5b.py` 是已验证的 NPU 入口。 + ## 模型权重 TeleFuser 不随软件包分发模型权重。[支持的模型](supported_models.md)页面会链接到各模型的 Cookbook， diff --git a/docs/zh/platforms.md b/docs/zh/platforms.md index 28d27424..b2986a7a 100644 --- a/docs/zh/platforms.md +++ b/docs/zh/platforms.md @@ -61,6 +61,19 @@ ROCm 支持面向使用 ROCm 7.x 与 PyTorch `+rocm` 构建的 AMD GPU；安装 ## NPU 与 CPU -NPU 平台通过 `torch_npu` 与 HCCL 分布式后端支持华为昇腾设备。平台层与算子分发层均已接入 NPU，但现有 -示例在 CUDA 上验证、部分示例在 ROCm 上验证 —— 生产使用前请先在目标 NPU 上完成验证。CPU 平台是未检测到 -加速器时的回退，面向测试以及显式请求 CPU 执行的 Pipeline。 +NPU 平台通过 `torch_npu` 与 HCCL 分布式后端支持华为昇腾设备。 + +- 注意力经原生回退路径使用 `TORCH_SDPA`；无需安装 `tf-kernel`、`flash_attn`、`sageattention` 或 `triton`。 +- 算子层在算子定义了 `forward_npu` 时优先选择，否则回退到 PyTorch 原生实现。目前尚未集成 NPU 优化内核， + Pipeline 完全运行在原生路径上。 +- 多卡推理通过 `hccl` 后端字符串使用 HCCL。并行 worker 队列经 CPU 中转张量（`torch_npu` 不提供可靠的 + 跨进程设备 IPC）；每个新起的 worker 组会分配独立的 `HCCL_IF_BASE_PORT`，避免并发组在 HCCL 数据面端口段 + 上冲突。 +- `torch.compile` 在 NPU 上未验证；NPU 示例以 eager 模式运行。 +- 已验证入口为 + [Wan2.2 TI2V-5B 文生视频](https://github.com/Tele-AI/TeleFuser/tree/main/examples/wan_video) + （`wan22_t2v_5b.py`）：示例自动检测平台、免修改运行，已在 Atlas 910B（CANN 8.2，torch 2.9 搭配版本匹配的 + `torch_npu`）上验证单卡以及 4 卡 CFG × Ulysses 并行（HCCL）。Wan2.2 A14B 复用同一代码路径，但尚未在 + NPU 硬件上运行；其他示例在生产使用前请先在目标 NPU 上完成验证。 + +CPU 平台是未检测到加速器时的回退，面向测试以及显式请求 CPU 执行的 Pipeline。 diff --git a/examples/wan_video/README.md b/examples/wan_video/README.md index e46cb48e..d497146b 100644 --- a/examples/wan_video/README.md +++ b/examples/wan_video/README.md @@ -35,6 +35,9 @@ Video generation using Wan2.1 and Wan2.2 models for Text-to-Video and Image-to-V - GPU: AMD ROCm GPUs for scripts ending in `_rocm.py`; validated on a Radeon RX 9070 (ROCm 7.2, `torch` built with `+rocm`). These examples use the PyTorch SDPA attention backend and need no tf-kernel, flash-attn, or SageAttention installation +- GPU: Huawei Ascend NPUs for `wan22_t2v_5b.py`, which auto-detects the platform; validated on an Atlas 910B + (CANN 8.2, torch 2.9 with a matching `torch_npu`). NPU execution uses the PyTorch SDPA attention backend and + needs no tf-kernel, flash-attn, SageAttention, or triton installation - Software: the standard TeleFuser installation; optional attention, FP8, Ray, and RIFE paths require their respective dependencies - Input assets: a readable image for I2V/FL2V and optional LoRA, distillation, cache, or RIFE weights for those variants @@ -463,6 +466,9 @@ python examples/wan_video/wan22_t2v_5b.py --resolution 480p --aspect_ratio 16:9 - CFG parallel enabled by default (cfg_scale=5.0) - Ulysses sequence parallelism for multi-GPU - 50-step UNPC sampling with sigma_shift=5.0 +- Platform auto-detection via `current_platform`: the script runs unmodified on CUDA and Ascend NPU hosts +- Validated on an Ascend Atlas 910B (CANN 8.2, torch 2.9 with a matching `torch_npu`): single-card and 4-card + CFG × Ulysses over HCCL, PyTorch SDPA attention, eager execution #### `wan22_14b_text_to_video_h100.py`