feat: add MindIE-SD as optional NPU attention and compilation backend (#1004)

Add MindIE-SD laser_attention and MindieSDBackend compile support for Ascend NPU,
with CLI-explicit control and zero environment variables.

Design:
- --attn _mindiesd_laser: use MindIE-SD laser attention
- --compile: auto-detect NPU + mindiesd -> torch.compile(backend=MindieSDBackend())
- Default attention is _native_npu; no env vars required (removed CACHE_DIT_ENABLE_MINDIESD_ATTN, CACHE_DIT_FORCE_DISABLE_MINDIESD_COMPILE_CONFIG)
- enable_cache() decoupled from compile; compile follows CUDA --compile pattern
- backend_selector.py moved from _utils/ to attention/ for semantic correctness
- Optional dependency: import mindiesd wrapped in try/except, graceful fallback
- NPU-only: all logic gated by torch.npu.is_available(), zero impact on CUDA/CPU

Files modified:
- attention/backends/npu.py: _mindiesd_laser_attention backend -> mindiesd.layers.attention_forward
- attention/backends/register.py: _MINDIESD_LASER enum
- attention/backend_selector.py: new BackendSelector.auto_select() -> _native_npu on NPU
- kernels/backend.py: KernelBackend.MINDIESD enum
- kernels/ops.py: MINDIESD kernel routing with PT fallback
- caching/cache_interface.py: auto-select attention backend in enable_cache()
- compile/utils.py: _maybe_apply_mindiesd_compile() -> torch.compile(backend=MindieSDBackend())
- _utils/utils.py: _mindiesd_laser in --attn CLI choices; _compile_transformer_module() detects NPU
- envs.py: removed MindIE-SD env vars

Verification (Ascend 910B, FLUX.1-dev 1024x1024, bf16, 28 steps):
  _native_npu (baseline):          489.7 ms/step
  _mindiesd_laser:                 479.3 ms/step  (-2.1%)
  _native_npu + compile:           447.3 ms/step  (-8.7%)
  _mindiesd_laser + compile:       422.3 ms/step  (-13.8%)

Co-authored-by: blian <lianbin@huawei.com>

Author: blian6

src/cache_dit/_utils/utils.py

@@ -7,7 +7,7 @@
 
 from ..logger import init_logger
 from ..platforms import current_platform
-from ..compile.utils import set_compile_configs
+from ..compile.utils import set_compile_configs, _maybe_apply_mindiesd_compile
 from ..distributed import ParallelismBackend, ParallelismConfig
 from ..caching import enable_cache, steps_mask
 from ..attention import set_attn_backend
@@ -698,6 +698,7 @@ def get_args(parse: bool = True, ) -> argparse.ArgumentParser | argparse.Namespa
       "sage",  # Need install sageattention: https://github.com/thu-ml/SageAttention
       "_native_npu",  # native npu attention
       "_npu_fia",  # npu fused infer attention
+      "_mindiesd_laser",  # MindIE-SD laser attention
     ],
   )
   # Ulysses context parallelism settings
@@ -1255,6 +1256,12 @@ def _compile_transformer_module(args, pipe, transformer, name):
     logger.warning(f"Cannot compile {name} module: {transformer_cls_name} Not a torch.nn.Module.")
     return transformer
 
+  # Auto-enable MindieSDBackend on NPU
+  compiled = _maybe_apply_mindiesd_compile(transformer, name, transformer_cls_name)
+  if compiled is not None:
+    setattr(pipe, name, compiled)
+    return compiled
+
   use_regional_compile = not args.disable_compile_repeated_blocks and hasattr(
     transformer, "compile_repeated_blocks")
 

src/cache_dit/attention/backend_selector.py

@@ -0,0 +1,33 @@
+import torch
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class BackendSelector:
+  _attn_backend: str | None = None
+  _selected: bool = False
+
+  @classmethod
+  def auto_select(cls, pipe_or_adapter) -> str | None:
+    if cls._selected:
+      return cls._attn_backend
+    device = cls._detect_device(pipe_or_adapter)
+    if device.type == "npu":
+      cls._attn_backend = "_native_npu"
+    cls._selected = True
+    return cls._attn_backend
+
+  @classmethod
+  def auto_select_kernel_backend(cls) -> str | None:
+    return None
+
+  @staticmethod
+  def _detect_device(pipe_or_adapter):
+    try:
+      if hasattr(pipe_or_adapter, "device"):
+        return pipe_or_adapter.device
+      param = next(pipe_or_adapter.parameters())
+      return param.device
+    except (StopIteration, AttributeError):
+      return torch.device("cpu")

src/cache_dit/attention/backends/npu.py

@@ -49,6 +49,8 @@ def _native_npu_attention(
   dropout_p: float = 0.0,
   scale: Optional[float] = None,
   return_lse: bool = False,
+  is_causal: bool = False,
+  enable_gqa: bool = False,
   _cp_config: Optional["_ContextParallelConfig"] = None,
 ) -> torch.Tensor:
   if return_lse:
@@ -97,6 +99,8 @@ def _npu_fused_infer_attention(
   dropout_p: float = 0.0,
   scale: Optional[float] = None,
   return_lse: bool = False,
+  is_causal: bool = False,
+  enable_gqa: bool = False,
   _cp_config: Optional["_ContextParallelConfig"] = None,
 ) -> torch.Tensor:
   if _cp_config is None:
@@ -128,3 +132,47 @@ def _npu_fused_infer_attention(
       _cp_config=_cp_config,
     )
   return out
+
+
+try:
+  from mindiesd.layers import attention_forward
+
+  _mindiesd_available = True
+except Exception:
+  _mindiesd_available = False
+  attention_forward = None
+
+if _mindiesd_available:
+
+  @_AttnBackendRegistry.register(
+    _AttnBackend._MINDIESD_LASER,
+    constraints=[],
+    supports_context_parallel=True,
+  )
+  def _mindiesd_laser_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    scale: Optional[float] = None,
+    return_lse: bool = False,
+    is_causal: bool = False,
+    enable_gqa: bool = False,
+    _cp_config: Optional["_ContextParallelConfig"] = None,
+  ) -> torch.Tensor:
+    if return_lse:
+      raise ValueError(
+        "MindIE-SD laser attention backend does not support setting `return_lse=True`.")
+    scale_val = scale if scale is not None else 1.0 / math.sqrt(query.shape[-1])
+    return attention_forward(
+      query,
+      key,
+      value,
+      attn_mask=attn_mask,
+      scale=scale_val,
+      fused=True,
+      head_first=False,
+      opt_mode="manual",
+      op_type="ascend_laser_attention",
+    )

src/cache_dit/attention/backends/register.py

@@ -38,6 +38,7 @@ class _AttnBackend(str, Enum):
   _SDPA_CUDNN = "_sdpa_cudnn"
   _NATIVE_NPU = "_native_npu"
   _NPU_FIA = "_npu_fia"
+  _MINDIESD_LASER = "_mindiesd_laser"
 
 
 def _default_active_backend() -> _AttnBackend:

src/cache_dit/caching/cache_interface.py

@@ -22,6 +22,16 @@
 logger = init_logger(__name__)
 
 
+def _auto_select_attention_backend(pipe_or_adapter) -> Optional[str]:
+  """Try to auto-select an optimal attention backend when none was specified."""
+  try:
+    from cache_dit.attention.backend_selector import BackendSelector
+
+    return BackendSelector.auto_select(pipe_or_adapter)
+  except Exception:
+    return None
+
+
 def enable_cache(
   pipe_or_adapter: Union[
     DiffusionPipeline,
@@ -352,6 +362,10 @@ def _enable_cache_impl(
     logger.warning("cache_config is None, skip cache acceleration for "
                    f"{pipe_or_adapter.__class__.__name__}.")
 
+  # Auto-select attention backend when none specified
+  if attention_backend is None and parallelism_config is None:
+    attention_backend = _auto_select_attention_backend(pipe_or_adapter)
+
   # Set custom attention backend for non-parallelism case
   if attention_backend is not None:
     if parallelism_config is not None:
@@ -457,6 +471,7 @@ def _enable_cache_impl(
             # Enable quantization for the specified component inplace
             quantized_component = quantize(component, quantize_config=config)
             setattr(pipe, name, quantized_component)
+
   return pipe_or_adapter
 
 

src/cache_dit/compile/utils.py

@@ -117,27 +117,33 @@ def set_compile_configs(
     # them to your needs and test the performance
     inductor_config.max_fusion_size = 64
     inductor_config.max_pointwise_cat_inputs = 8
-    inductor_config.triton.cudagraphs = cuda_graphs
-    inductor_config.triton.use_block_ptr = False
-    inductor_config.triton.codegen_upcast_to_fp32 = True
-
-    # Copy from https://pytorch.org/blog/accelerating-generative-ai-3/
-    inductor_config.conv_1x1_as_mm = True
-    inductor_config.coordinate_descent_tuning = True
-    inductor_config.coordinate_descent_check_all_directions = True
-    inductor_config.epilogue_fusion = False
-
-    # Enable epilogue and prologue fusion
-    if ENV.CACHE_DIT_EPILOGUE_PROLOGUE_FUSION or kwargs.get(
-        "epilogue_prologue_fusion",
-        False,
-    ):
-      inductor_config.epilogue_fusion = True
-      inductor_config.prologue_fusion = True
-      inductor_config.epilogue_fusion_first = True
 
-    # Dead code elimination
-    inductor_config.dce = True  # default is False
+    if current_platform.device_type == "npu":
+      # NPU: skip CUDA-specific inductor configs (triton, coordinate_descent, etc)
+      inductor_config.dce = True  # default is False
+      inductor_config.epilogue_fusion = False
+    else:
+      inductor_config.triton.cudagraphs = cuda_graphs
+      inductor_config.triton.use_block_ptr = False
+      inductor_config.triton.codegen_upcast_to_fp32 = True
+
+      # Copy from https://pytorch.org/blog/accelerating-generative-ai-3/
+      inductor_config.conv_1x1_as_mm = True
+      inductor_config.coordinate_descent_tuning = True
+      inductor_config.coordinate_descent_check_all_directions = True
+      inductor_config.epilogue_fusion = False
+
+      # Enable epilogue and prologue fusion
+      if ENV.CACHE_DIT_EPILOGUE_PROLOGUE_FUSION or kwargs.get(
+          "epilogue_prologue_fusion",
+          False,
+      ):
+        inductor_config.epilogue_fusion = True
+        inductor_config.prologue_fusion = True
+        inductor_config.epilogue_fusion_first = True
+
+      # Dead code elimination
+      inductor_config.dce = True  # default is False
 
     # May need to force disable all cache
     if force_disable_compile_caches:
@@ -153,3 +159,19 @@ def set_compile_configs(
       inductor_config.cuda.use_fast_math = use_fast_math
   except Exception:
     pass
+
+
+def _maybe_apply_mindiesd_compile(module, module_name, module_cls_name):
+  # Auto-apply MindieSDBackend compile on NPU when mindiesd is available.
+  # Returns the compiled module if compiled, None if MindIE-SD not applicable.
+  try:
+    import mindiesd  # noqa F401
+
+    if not hasattr(torch, 'npu') or not torch.npu.is_available():
+      return None
+    from mindiesd.compilation import MindieSDBackend
+
+    logger.info(f"Compiling {module_name}: {module_cls_name} with MindieSDBackend ...")
+    return torch.compile(module, backend=MindieSDBackend(), dynamic=True)
+  except Exception:
+    return None

src/cache_dit/kernels/backend.py

@@ -7,6 +7,7 @@ class KernelBackend(Enum):
   TRITON = "Triton"
   CUDA = "CUDA"
   CUTEDSL = "CuteDSL"
+  MINDIESD = "MindIESD"
   NONE = "None"
 
   @classmethod
@@ -41,4 +42,11 @@ def is_supported(cls, backend: "KernelBackend") -> bool:
         return True
       except ImportError:
         return False
+    if backend == cls.MINDIESD:
+      try:
+        import mindiesd  # noqa F401
+
+        return True
+      except Exception:
+        return False
     return False

src/cache_dit/kernels/ops.py

@@ -14,6 +14,10 @@ def _select_cuda_backend() -> KernelBackend:
   return KernelBackend.CUDA
 
 
+def _select_mindiesd_backend() -> KernelBackend:
+  return KernelBackend.MINDIESD
+
+
 def _select_kernel_backend() -> KernelBackend:
   """Select the default backend for kernel ops.
 
@@ -149,6 +153,20 @@ def _fused_merge_attn_states_impl(
       suff_out,
       suff_lse,
     )
+  if backend == KernelBackend.MINDIESD:
+    max_lse = torch.max(prev_lse, suff_lse)
+    prev_scale = torch.exp(prev_lse - max_lse)
+    suff_scale = torch.exp(suff_lse - max_lse)
+    denom = prev_scale + suff_scale
+    denom = denom + 1e-12
+    out = torch.where(
+      denom > 1e-12,
+      (prev_out * prev_scale.unsqueeze(-1) + suff_out * suff_scale.unsqueeze(-1)) /
+      denom.unsqueeze(-1),
+      prev_out,
+    )
+    lse = max_lse + torch.log(denom)
+    return out, lse
   else:
     raise ValueError(_ERROR_TEMPLATE.format(backend))