Refactor/npu transformer cleanup

.gitignore

@@ -13,3 +13,11 @@ dist/
 CLAUDE.md
 .claude/
 .kiro/
+
+# Performance benchmarks & results (internal use only)
+benchmarks/
+results/
+
+# PR artifacts
+PR_BODY.md
+.pr272.diff

diffsynth_engine/args.py

@@ -117,6 +117,11 @@ def parse_cli_args() -> Dict[str, Any]:
         action="store_true",
         help="Compile repeated transformer blocks with torch.compile",
     )
+    optimization_group.add_argument(
+        "--compile-ffn",
+        action="store_true",
+        help="Compile only FFN (MLP) blocks with torch.compile (finer-grained than --use-torch-compile)",
+    )
 
     # Parallelism configuration group
     parallel_group = parser.add_argument_group("Parallelism Configuration")
@@ -184,6 +189,7 @@ def parse_cli_args() -> Dict[str, Any]:
 
     # Optimization configuration
     args_dict["use_torch_compile"] = args.use_torch_compile
+    args_dict["compile_ffn"] = args.compile_ffn
 
     # Parallelism configuration
     args_dict["parallelism"] = args.parallelism

diffsynth_engine/configs/base.py

@@ -6,6 +6,7 @@
 from diffsynth_engine.layers.attention import AttentionType
 from diffsynth_engine.registry import get_attn_backend
 from diffsynth_engine.utils import logging
+from diffsynth_engine.platforms import get_device_type, resolve_platform
 
 logger = logging.get_logger(__name__)
 
@@ -27,7 +28,7 @@ class PipelineConfig:
     model_dtype: torch.dtype = torch.bfloat16
     text_encoder_dtype: torch.dtype = torch.bfloat16
     vae_dtype: torch.dtype = torch.float32
-    device: str | torch.device = "cuda"
+    device: str | torch.device = "auto"
 
     pipeline_class_name: str | None = None
 
@@ -42,6 +43,7 @@ class PipelineConfig:
 
     # optimization
     use_torch_compile: bool = False
+    compile_ffn: bool = False
 
     # parallelism
     parallelism: int = 1
@@ -62,6 +64,7 @@ def __post_init__(self):
         self.attn_type = str(self.attn_type)
         init_parallel_config(self)
         validate_attn_config(self)
+        init_device_config(self)
 
 
 def init_parallel_config(config: PipelineConfig):
@@ -109,3 +112,11 @@ def validate_attn_config(config: PipelineConfig):
     if config.sp_ring_degree is not None and config.sp_ring_degree > 1:
         if not attn_backend.supports_ring_attention():
             raise ValueError(f"Attention backend {config.attn_type!r} does not support ring attention.")
+
+
+def init_device_config(config: PipelineConfig):
+    if config.device is None or (isinstance(config.device, str) and config.device.lower() in ("auto", "")):
+        config.device = get_device_type()
+        return
+    # Validate that the explicit device type is registered (fail fast at construction).
+    resolve_platform(config.device)

diffsynth_engine/distributed/parallel_state.py

@@ -13,7 +13,6 @@
 
 import torch
 import torch.distributed
-from torch.cuda import device_count, set_device
 
 from diffsynth_engine.distributed.group_coordinator import (
     GroupCoordinator,
@@ -22,7 +21,11 @@
 )
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import IDLE_TIMEOUT_SEC
-from diffsynth_engine.utils.platform import get_torch_distributed_backend
+from diffsynth_engine.utils.platform import (
+    device_count,
+    get_torch_distributed_backend,
+    set_device,
+)
 
 logger = logging.get_logger(__name__)
 
@@ -425,28 +428,30 @@ def init_distributed_environment(
         distributed_init_method,
         backend,
     )
+    # local_rank is not available in torch ProcessGroup,
+    # see https://github.com/pytorch/pytorch/issues/122816
+    if local_rank == -1:
+        # local rank not set, this usually happens in single-node
+        # setting, where we can use rank as local rank
+        if distributed_init_method == "env://":
+            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+        else:
+            local_rank = rank if rank >= 0 else 0
+
     if not torch.distributed.is_initialized():
         assert distributed_init_method is not None, (
             "distributed_init_method must be provided when initializing distributed environment"
         )
+        # Bind device before init_process_group (required by HCCL on Ascend).
+        set_device(local_rank % max(device_count(), 1))
         # this backend is used for WORLD
         torch.distributed.init_process_group(
             backend=backend,
             init_method=distributed_init_method,
             world_size=world_size,
             rank=rank,
         )
-        set_device(torch.distributed.get_rank() % device_count())
-    # set the local rank
-    # local_rank is not available in torch ProcessGroup,
-    # see https://github.com/pytorch/pytorch/issues/122816
-    if local_rank == -1:
-        # local rank not set, this usually happens in single-node
-        # setting, where we can use rank as local rank
-        if distributed_init_method == "env://":
-            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
-        else:
-            local_rank = rank
+        set_device(torch.distributed.get_rank() % max(device_count(), 1))
     global _WORLD
     if _WORLD is None:
         ranks = list(range(torch.distributed.get_world_size()))

diffsynth_engine/engine.py

@@ -1,14 +1,14 @@
 from typing import Any
 
 import torch.multiprocessing as mp
-from torch.cuda import set_device
 
 from diffsynth_engine.configs import PipelineConfig
 from diffsynth_engine.registry import (
     get_pipeline_class,
     get_pipeline_class_name,
 )
 from diffsynth_engine.utils import logging
+from diffsynth_engine.utils.platform import set_device
 from diffsynth_engine.utils.torch_profiler import TorchProfiler
 from diffsynth_engine.worker import run_worker_loop
 

diffsynth_engine/layers/__init__.py

@@ -0,0 +1,5 @@
+from diffsynth_engine.layers.transformer_helper import RMSNorm
+
+__all__ = [
+    "RMSNorm",
+]

diffsynth_engine/layers/attention/__init__.py

@@ -1,9 +1,13 @@
 from .backends.abstract import AttentionMetadata, AttentionType
+from .factory import create_parallel_attention
 from .layer import LocalAttention, USPAttention
+from .ascend_long_context import AscendLongContextAttention
 
 __all__ = [
     "AttentionType",
     "AttentionMetadata",
     "LocalAttention",
     "USPAttention",
+    "AscendLongContextAttention",
+    "create_parallel_attention",
 ]