编译特性

MindIE SD 基于 PyTorch 的 torch.compile 编译器提供自定义后端 MindieSDBackend()，在昇腾芯片上提供两套互补的加速能力：

• Pattern 融合：利用 Pattern Matcher 将常见算子组合自动替换为昇腾融合算子，减少 kernel 启动开销
• ACLGraph 加速：通过 torch.npu.NPUGraph 将计算图捕获为静态执行图，replay 时跳过动态图调度

两项能力通过 CompilationConfig 统一控制。

基本用法

两种加速能力共享相同的入口：对模型或其子模块调用 torch.compile 并指定 MindieSDBackend()。

对 transformer 整体 compile：

pipe = FluxPipeline.from_pretrained(...)
transformer = torch.compile(pipe.transformer, backend=MindieSDBackend())
setattr(pipe, "transformer", transformer)

对单个 Module 使用装饰器：

@torch.compile(backend=MindieSDBackend())
class FluxSingleTransformerBlock(nn.Module):

对 forward 函数使用装饰器：

class FluxSingleTransformerBlock(nn.Module):
    @torch.compile(backend=MindieSDBackend())
    def forward(...):

Pattern 融合

MindieSDBackend() 内置了多组算子融合 Pattern，编译时自动匹配并替换为昇腾优化算子。各 Pattern 的开关可通过 CompilationConfig.fusion_patterns 单独控制：

from mindiesd.compilation import CompilationConfig

CompilationConfig.fusion_patterns.enable_rms_norm = False      # 关闭 RMSNorm 融合
CompilationConfig.fusion_patterns.enable_rope = False          # 关闭 RoPE 融合
CompilationConfig.fusion_patterns.enable_adalayernorm = False  # 关闭 adaLN 融合
CompilationConfig.fusion_patterns.enable_fast_gelu = False     # 关闭 fastGelu 融合
CompilationConfig.fusion_patterns.enable_mul_add = False       # 关闭 Mul+Add 融合

各融合算子的详细 API 说明请参见 core_layers.md 中的融合算子章节。

接口说明

MindieSDBackend

from mindiesd.compilation import MindieSDBackend

作为 torch.compile 的 backend 参数传入，自动启用 Pattern 融合和 ACLGraph 加速。

CompilationConfig

from mindiesd.compilation import CompilationConfig

Pattern 融合的开关通过 fusion_patterns 控制：

配置项	默认值	说明
fusion_patterns.enable_rms_norm	True	RMSNorm 融合
fusion_patterns.enable_rope	True	RoPE 融合
fusion_patterns.enable_adalayernorm	True	AdaLayerNorm 融合
fusion_patterns.enable_fast_gelu	True	fastGELU 融合
fusion_patterns.enable_mul_add	True	Mul+Add 融合

ACLGraph 加速

在 Pattern 融合的基础上，可进一步启用 ACLGraph 将优化后的图捕获为静态执行计划。

配置方式

aclgraph_only 与 aclgraph_with_compile 互斥，同时开启时 aclgraph_with_compile 优先级更高。

配置项	默认值	说明
aclgraph_only	False	仅 ACLGraph，跳过 Pattern 融合
aclgraph_with_compile	False	先 Pattern 融合，再捕获为 ACLGraph
enable_freezing	True	编译前是否执行常量折叠
safe_output_mode	True	ACLGraph replay 时是否 clone 输出
graph_log_url	None	调试用 graph transform 日志 URL

使用示例

在上方 基本用法 的基础上，调用前配置 CompilationConfig 即可启用：

from mindiesd.compilation import CompilationConfig

CompilationConfig.aclgraph_with_compile = True
# 之后调用 torch.compile(..., backend=MindieSDBackend()) 即自动启用

变长输入处理

语音等场景下输入长度不固定，可通过外部 padding 适配：

max_len = 512

model = torch.compile(transformer, backend=MindieSDBackend())
_ = model(torch.randn(max_len, dim, device="npu"))  # 触发捕获

for audio_chunk in chunks:
    actual_len = audio_chunk.shape[0]
    padded = torch.nn.functional.pad(audio_chunk, (0, 0, 0, max_len - actual_len))
    output = model(padded)[:actual_len]

限制与注意事项

• 环境依赖：仅昇腾 NPU 环境支持
• 输入 shape：运行时输入 shape 须与捕获时一致，变更会触发重新捕获
• 动态特性：不支持动态 shape、动态 control flow 或 conditional branching
• 首次耗时：首次触发图捕获存在一次性耗时开销
• graph.update：不提供 graph.update 接口（该接口用于 LLM 场景动态注入 attention metadata，SD 场景不需要）
• 配置时机：CompilationConfig 需在 torch.compile() 调用前完成配置

问题定位技巧

• 相关的定位手段与PyTorch的compile一致，mindie_sd_backend.py中定义了日志模块，开启后，可以观察到pattern使能前后的图变化情况。配合torch.compile缩小范围，可以识别pattern失效的原因。
• 通过控制compile的范围，可以有效控制问题定位的范围。
• 其他定位手段可以参考PyTorch官网。