[FEATURE] Run Wan2.2 pipelines on Ascend NPU

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()

telefuser/distributed/device_mesh.py

@@ -13,22 +13,25 @@
 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
     For USP (Unified Sequence Parallelism), ring and ulysses form a 2D sub-mesh.
 
     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

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

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,

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

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)
 
 

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(

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