整体说明

TorchAir是为Ascend Extension for PyTorch提供图模式能力的扩展库，支持用户使用PyTorch在昇腾设备上进行图模式推理。本章详细介绍了TorchAir框架提供的Python API，具体参见表1。

表1 接口列表
接口分类	接口名	接口说明
torchair	CompilerConfig类	该类用于构造传入torch.compiler backend的配置。
	dynamo_export	用于导出由TorchAir生成的离线图（air格式）。
	get_compiler	获取能够在NPU上运行的图编译器。
	get_npu_backend	获取能够在NPU上运行的图编译后端npu_backend，可作为backend参数传入torch.compile。
	use_internal_format_weight	将模型中的权重weight转成TorchAir定义的内部私有格式。
	register_fx_node_ge_converter	将自定义算子注册到TorchAir框架中。
	patch_for_hcom	针对PyTorch 2.1版本中不支持入图的集合通信算子提供的补丁函数，实现部分集合通信算子入图。
	register_replacement	将自定义算子融合规则注册到TorchAir中，在FX图编译后对图进行算子融合优化。
torchair.ge	DataType类	数据类型的枚举类，提供了GE的data type定义，方便实现converter函数时调用。
	Format类	数据格式的枚举类，提供了GE的data format定义，方便实现converter函数时调用。
	Tensor类	提供Tensor定义，用于算子入图的converter函数入参类型声明。
	TensorSpec类	提供TensorSpec定义，表示算子在Meta推导过程中得到的性能，当前用于算子入图的converter函数入参类型声明。
	Const	算子converter中的构图元素，表示一个Const节点，即图中的常量值。
	Cast	算子converter中的构图元素，表示一个Cast节点，即图中Tensor的类型转换方法。
	Clone	算子Converter中的构图元素，表示一个Clone节点，该节点可实现图上任意单个Tensor的拷贝。
	custom_op	基于算子原型（IR）实现算子converter函数，完成PyTorch IR与GE IR的转换，方便自定义算子入图。
torchair.inference	cache_compile	开启模型编译缓存功能（aclgraph）时需要调用该接口实现模型编译缓存。
	readable_cache	开启模型编译缓存功能（aclgraph）时需要调用该接口读取封装的func函数缓存文件compiled_module，并以可读文件（格式不限，如py、txt）格式呈现。
	set_dim_gears	开启动态shape图分档执行功能时需要调用该接口设置图被划分的档位。
torchair.ops	npu_print	图执行过程中，打印执行脚本中目标tensor值。
	npu_fused_infer_attention_score	图模式场景下的FlashAttention算子，既可以支持全量计算场景（PromptFlashAttention），也可支持增量计算场景（IncreFlashAttention）。
	npu_fused_infer_attention_score_v2	图模式场景下的增强版FlashAttention算子，支持全量和增量计算场景。
	npu_create_tagged_event	根据tag创建一个唯一的事件对象torch.npu.Event（torch.cuda.Event的NPU形式）用于协调NPU上不同stream之间的同步操作。
	npu_tagged_event_record	与torch.npu.Event.record方法类似（torch.cuda.Event.record的NPU形式），用于记录当前流中的事件。
	npu_tagged_event_wait	与torch.npu.Event.wait方法类似（torch.cuda.Event.wait的NPU形式），用于让当前流等待指定的事件完成。
	npu_record_tagged_stream	与PyTorch原生torch.Tensor.record_stream接口类似，用于确保张量在特定NPU流完成之前不会被释放。
	record	torchair.ops.record用于显式地在当前Stream上下发一个任务，其返回值可以被torchair.ops.wait等待。
	wait	用于在多流间控制时序关系，torchair.ops.wait表示当前流需要在传入的tensor对应的节点执行结束后，再继续执行。
torchair.scope	npu_stream_switch	图执行过程中，指定图内多个算子分发到不同stream做并行计算。
	npu_wait_tensor	图执行过程中，控制图内多stream并行计算时序。
	super_kernel	图执行过程中，标记图内能融合为SuperKernel的上下文算子范围。
	limit_core_num	图执行过程中，指定图内范围内的算子执行时的AI Core数和Vector Core数。
	op_never_timeout	针对GE图中的算子添加_op_exec_never_timeout属性，即配置算子不超时，使其不参与超时检测。
torchair.llm_datadist	create_npu_tensors	通过一串Device地址创建PyTorch在NPU上的Tensors。主要用于创建大模型中的KV Cache Tensors，所有KV Cache的shape和dtype都一致。

父主题： API参考