已开启

1 #1

rmch创建于 2025年10月27日

rmch成员

2025年10月27日创建

import os
os.environ['TORCHINDUCTOR_MAX_AUTOTUNE']='1'
import torch
from torch._inductor.compile_fx import clone_preserve_strides
from torch._dynamo.testing import rand_strided
from torch import device

import torch_npu
from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as npu_config

from torch import empty_strided, empty, randn
from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.codecache import CustomAsyncCompile
from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.utils import (
logger,
)
import logging
logger.setLevel(logging.INFO)
async_compile = CustomAsyncCompile()

file_path = os.path.abspath(file)
dir_path = os.path.dirname(file_path)

mlir_fused__npu_dtype_cast_backward__to_copy_add_div_mul_pow_sum_19 = async_compile.mlir_auto_fallback('mlir_fused__npu_dtype_cast_backward__to_copy_add_div_mul_pow_sum_19', '''
module {
func.func @mlir_fused__npu_dtype_cast_backward__to_copy_add_div_mul_pow_sum_19(%arg0: tensor<3643x3584xbf16>, %arg1: tensor<3643x3584xbf16>, %arg2: tensor<3584xbf16>, %arg3: tensor<3643x3584xf32>, %arg4: tensor<3643x1xf32>, %arg5: tensor<3643x3584xbf16>) -> (tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) attributes {hacc.entry, hacc.function_kind = #hacc.function_kind} {
%c3_i64 = arith.constant 3 : i64
%c3584_i64 = arith.constant 3584 : i64
%c1_i64 = arith.constant 1 : i64
%cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant -5.000000e-01 : f32
%cst_1 = arith.constant 2.000000e+00 : f32
%0 = tensor.empty() : tensor<3643x3584xi64>
%1 = linalg.fill ins(%c1_i64 : i64) outs(%0 : tensor<3643x3584xi64>) -> tensor<3643x3584xi64>
%2 = tensor.empty() : tensor<3643x3584xbf16>
%3 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%1 : tensor<3643x3584xi64>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%4 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%arg1, %3 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%5 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%arg0, %4 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%broadcasted = linalg.broadcast ins(%arg2 : tensor<3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) dimensions = [0]
%6 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%5, %broadcasted : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%7 = tensor.empty() : tensor<3643x3584xf32>
%8 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%6 : tensor<3643x3584xbf16>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%9 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%8, %arg3 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%10 = tensor.empty() : tensor<3643xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<3643xf32>) -> tensor<3643xf32>
%reduced = linalg.reduce ins(%9 : tensor<3643x3584xf32>) outs(%11 : tensor<3643xf32>) dimensions = [1]
(%in: f32, %init: f32) {
%36 = arith.addf %in, %init : f32
linalg.yield %36 : f32
}
%expanded = tensor.expand_shape %reduced [[0, 1]] output_shape [3643, 1] : tensor<3643xf32> into tensor<3643x1xf32>
%collapsed = tensor.collapse_shape %arg4 [[0, 1]] : tensor<3643x1xf32> into tensor<3643xf32>
%broadcasted_2 = linalg.broadcast ins(%collapsed : tensor<3643xf32>) outs(%7 : tensor<3643x3584xf32>) dimensions = [1]
%12 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%arg3, %broadcasted_2 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%13 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%12 : tensor<3643x3584xf32>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%14 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%5, %13 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%15 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%8, %broadcasted_2 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%16 = tensor.empty() : tensor<3643x1xf32>
%17 = linalg.fill ins(%cst_0 : f32) outs(%16 : tensor<3643x1xf32>) -> tensor<3643x1xf32>
%18 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%expanded, %17 : tensor<3643x1xf32>, tensor<3643x1xf32>) outs(%16 : tensor<3643x1xf32>) -> tensor<3643x1xf32>
%19 = tensor.empty() : tensor<3643x1xi64>
%20 = linalg.fill ins(%c3_i64 : i64) outs(%19 : tensor<3643x1xi64>) -> tensor<3643x1xi64>
%21 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%20 : tensor<3643x1xi64>) outs(%16 : tensor<3643x1xf32>) -> tensor<3643x1xf32>
%22 = hfusion.elemwise_binary {fun = #hfusion.binary_fn} ins(%arg4, %21 : tensor<3643x1xf32>, tensor<3643x1xf32>) outs(%16 : tensor<3643x1xf32>) -> tensor<3643x1xf32>
%23 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%18, %22 : tensor<3643x1xf32>, tensor<3643x1xf32>) outs(%16 : tensor<3643x1xf32>) -> tensor<3643x1xf32>
%collapsed_3 = tensor.collapse_shape %23 [[0, 1]] : tensor<3643x1xf32> into tensor<3643xf32>
%broadcasted_4 = linalg.broadcast ins(%collapsed_3 : tensor<3643xf32>) outs(%7 : tensor<3643x3584xf32>) dimensions = [1]
%24 = linalg.fill ins(%c3584_i64 : i64) outs(%0 : tensor<3643x3584xi64>) -> tensor<3643x3584xi64>
%25 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%24 : tensor<3643x3584xi64>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%26 = linalg.elemwise_binary {fun = #linalg.binary_fn

} ins(%broadcasted_4, %25 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%27 = linalg.fill ins(%cst_1 : f32) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%28 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%arg3, %27 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%29 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%26, %28 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%30 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%1 : tensor<3643x3584xi64>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%31 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%29, %30 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%32 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%15, %31 : tensor<3643x3584xf32>, tensor<3643x3584xf32>) outs(%7 : tensor<3643x3584xf32>) -> tensor<3643x3584xf32>
%33 = hfusion.cast {round_mode = #hfusion.round_mode} ins(%32 : tensor<3643x3584xf32>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%34 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%33, %3 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
%35 = linalg.elemwise_binary {fun = #linalg.binary_fn} ins(%arg5, %34 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>) outs(%2 : tensor<3643x3584xbf16>) -> tensor<3643x3584xbf16>
return %14, %35 : tensor<3643x3584xbf16>, tensor<3643x3584xbf16>
}
}
''', kernel_meta={'device_str': 'npu', 'device_index': 0, 'num_outputs': 2, 'non_contiguous_indices': {'inputs': [], 'outputs': []}, 'dynamic': False, 'mutated_indices': [], 'traced_graph_cache': '/tmp/torchinductor_root/traced_graph_cache', 'num_call_functions': 2, 'traced_graph_hash': 'cyn3ghbjhlhnphdvtzn7jtpyaklt3wpsi3lmbhqsnhrg26rgkrrl', 'signature': {0: '*bf16', 1: '*bf16', 2: '*bf16', 3: '*fp32', 4: '*fp32', 5: '*bf16', 6: '*bf16', 7: '*bf16'}, 'ranks': [2, 2, 1, 2, 2, 2, 2, 2], 'kernel_hash': 'b680bbb8623fc96ba61fb272c21f6ec64f5c4745c23023ee81a18abddd29a2a4'})

from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import get_current_raw_stream as get_raw_stream

async_compile.wait(globals())
del async_compile

stream0 = get_raw_stream(0)

args = [torch.randn(torch.Size([3643, 3584]), device="npu", dtype=torch.bfloat16),
torch.randn(torch.Size([3643, 3584]), device="npu", dtype=torch.bfloat16),
torch.randn(torch.Size([3584]), device="npu", dtype=torch.bfloat16),
torch.randn(torch.Size([3643, 3584]), device="npu", dtype=torch.float32),
torch.randn(torch.Size([3643, 1]), device="npu", dtype=torch.float32),
torch.randn(torch.Size([3643, 3584]), device="npu", dtype=torch.bfloat16),
torch.empty(torch.Size([3643, 3584]), device="npu", dtype=torch.bfloat16),
torch.empty(torch.Size([3643, 3584]), device="npu", dtype=torch.bfloat16),
]
args = [arg.npu() if isinstance(arg, torch.Tensor) else arg for arg in args]
num_args = 2

mlir_fused__npu_dtype_cast_backward__to_copy_add_div_mul_pow_sum_19.run(
*args,
stream=stream0)

output1 = args[num_args:]

zichun_ye

3月26日评论：

https://pytorch-package.obs.cn-north-4.myhuaweicloud.com/pta/PR/32297/torch_npu_x86_64.tar.gz

rmch成员

5月28日评论：

Guard Filter（Guard 过滤）

简介

PyTorch Dynamo（。。。。）在每次进入编译入口时会生成一组 guard（守卫条件），用于检测运行时状态是否发生变化。当 guard 失效时会触发重编译（recompilation），严重影响推理性能。

使用场景

在 NPU 推理场景中，以下几类 guard 频繁触发重编译，但实际运行时对应的状态并不会影响编译产物的正确性：

Guard 类型	触发场景
`DICT_VERSION` / `DICT_KEYS` / `DICT_KEYS_MATCH` / `DICT_CONTAINS`	HuggingFace `generation_config`、KV cache state dict、attention kwargs，每步 version 自增
`TYPE_MATCH` / `OPTIONAL_TENSOR`	`past_key_values=None`（prefill）↔ `tuple`（decode）、`attention_mask=None` ↔ Tensor
`HASATTR` / `NOT_PRESENT_IN_GENERIC_DICT`	特性开关、首次 forward 后才挂上的字段
`GRAD_MODE` / `TORCH_FUNCTION_STATE` / `DEFAULT_DEVICE` / `DETERMINISTIC_ALGORITHMS` / `AUTOCAST_STATE` / `FSDP_TRAINING_STATE`	进程级一次性配置，但每个编译入口都会生成 guard

PyTorch 原生提供了 guard_filter_fn 编译选项，允许用户自定义过滤逻辑，选择性跳过不必要的 guard。

Caution

guard_filter_fn具有unsafe语义，如果在运行时改变了被过滤guard对应的状态，编译产物将静默地产生错误结果。

使用前请确认：

被过滤的状态在整个推理过程中确实不会发生语义上的变化。
通过 parity 测试验证输出正确性。
通过 recompile count 测试验证重编译已消除。

使用指导

接口原型：

torch.compile(model, options={"guard_filter_fn": filter_fn})

参数说明：

filter_fn 签名：

def filter_fn(entries: list) -> list[bool]:
    """
    参数:
        entries: guard 条目列表，每个条目包含以下属性：
            - guard_type (str): guard 类型，如 "DICT_VERSION"、"GRAD_MODE" 等
            - name (str): guard 关联的变量名
            - is_global (bool): 是否为全局变量的 guard
            - value: guard 关联的值（可选）
    返回:
        与 entries 等长的布尔列表，True 表示保留该 guard，False 表示过滤掉
    """

[!NOTICE]

PyTorch还提供了内置的unsafe helper：

torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe：跳过内置 nn.Module 属性变化的 guard

torch.compiler.skip_guard_on_all_nn_modules_unsafe：跳过所有 nn.Module 属性变化的 guard

torch.compiler.skip_guard_on_globals_unsafe：跳过全局变量的 guard

使用示例

示例 1：过滤字典版本 guard

适用于 HuggingFace 模型推理中 generation_config 等字典频繁变化的场景：

import torch
import torch_npu

_DICT_GUARD_TYPES = frozenset({
    "DICT_VERSION", "DICT_KEYS", "DICT_KEYS_MATCH", "DICT_CONTAINS",
})

def filter_dict_guards(entries):
    return [entry.guard_type not in _DICT_GUARD_TYPES for entry in entries]

model = MyModel().npu()
compiled = torch.compile(model, options={"guard_filter_fn": filter_dict_guards})

示例 2：过滤运行时状态 guard

适用于推理管线中交替切换 torch.no_grad() / torch.enable_grad() 的场景：

import torch
import torch_npu

_RUNTIME_STATE_GUARD_TYPES = frozenset({
    "GRAD_MODE", "TORCH_FUNCTION_STATE", "GLOBAL_STATE",
    "DEFAULT_DEVICE", "DETERMINISTIC_ALGORITHMS", "AUTOCAST_STATE",
    "FSDP_TRAINING_STATE",
})

def filter_runtime_state_guards(entries):
    return [entry.guard_type not in _RUNTIME_STATE_GUARD_TYPES for entry in entries]

model = MyModel().npu()
compiled = torch.compile(model, options={"guard_filter_fn": filter_runtime_state_guards})

# 切换 grad mode 不会触发重编译
with torch.no_grad():
    compiled(x)
with torch.enable_grad():
    compiled(x)  # 不会重编译

示例 3：组合多类 guard 过滤

一次性过滤多类不必要的 guard：

import torch
import torch_npu

_FILTER_GUARD_TYPES = frozenset({
    # 字典版本
    "DICT_VERSION", "DICT_KEYS", "DICT_KEYS_MATCH", "DICT_CONTAINS",
    # 可选类型
    "TYPE_MATCH", "OPTIONAL_TENSOR",
    # hasattr
    "HASATTR", "NOT_PRESENT_IN_GENERIC_DICT",
    # 运行时状态
    "GRAD_MODE", "TORCH_FUNCTION_STATE", "GLOBAL_STATE",
    "DEFAULT_DEVICE", "DETERMINISTIC_ALGORITHMS", "AUTOCAST_STATE",
    "FSDP_TRAINING_STATE",
})

def npu_guard_filter(entries):
    return [entry.guard_type not in _FILTER_GUARD_TYPES for entry in entries]

model = MyModel().npu()
compiled = torch.compile(model, options={"guard_filter_fn": npu_guard_filter})

示例 4：按变量名或属性过滤

精确过滤特定变量的 guard：

import torch
import torch_npu

def filter_by_name(entries):
    return [
        not (entry.name == "y" and entry.value is None)
        for entry in entries
    ]

@torch.compile(fullgraph=True, options={"guard_filter_fn": filter_by_name})
def fn(x, y):
    if y is not None:
        x += y
    return x

示例 5：过滤全局变量 guard

import torch
import torch_npu

def filter_globals(entries):
    return [not entry.is_global for entry in entries]

model = MyModel().npu()
compiled = torch.compile(model, options={"guard_filter_fn": filter_globals})

示例 6：结合内置 helper 使用

import torch
import torch_npu

# 使用 PyTorch 内置的 nn.Module guard 过滤
model = MyModel().npu()
compiled = torch.compile(
    model,
    options={
        "guard_filter_fn": torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe
    },
)

调试与验证

确认重编译是否消除

使用 torch.compiler.set_stance("fail_on_recompile") 验证：

compiled = torch.compile(model, options={"guard_filter_fn": npu_guard_filter})

# 首次编译
compiled(x)

# 验证不会重编译
with torch.compiler.set_stance("fail_on_recompile"):
    compiled(x)  # 如果重编译会抛出异常

查看 guard 日志

通过环境变量开启 guard 日志，定位触发重编译的 guard 类型：

TORCH_LOGS=guards,recompiles python your_script.py

验证输出正确性

过滤 guard 后务必验证编译输出与 eager 模式一致：

model.eval()
x = torch.randn(2, 8).npu()

with torch.no_grad():
    eager_out = model(x)

compiled = torch.compile(model, options={"guard_filter_fn": npu_guard_filter})
with torch.no_grad():
    compiled_out = compiled(x)

assert torch.allclose(eager_out, compiled_out, atol=1e-5)

约束说明

PyTorch版本：必须为2.4或更高版本。
touch_npu版本：必须安装与PyTorch版本配套的版本，具体请参考版本说明