整体说明

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

表1 接口列表
接口分类	接口名	接口说明
torchair	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版本中不支持入图的集合通信算子提供的补丁函数，实现部分集合通信算子入图。
torchair.ge	Const	算子converter中的构图元素，表示一个Const节点，即图中的常量值。
	Cast	算子converter中的构图元素，表示一个Cast节点，即图中Tensor的类型转换方法。
	custom_op	基于算子原型（IR）实现算子converter函数，完成PyTorch IR与GE IR的转换，方便自定义算子入图。
torchair.inference	cache_compile	开启模型编译缓存功能时需要调用该接口实现模型编译缓存。
	readable_cache	开启模型编译缓存功能时需要调用该接口读取封装的func函数缓存文件compiled_module，并以可读文件（格式不限，如py、txt）格式呈现。
	set_dim_gears	开启动态shape图分档执行功能时需要调用该接口设置图被划分的档位。
torchair.ops	npu_print	图执行过程中，打印执行脚本中目标Tensor值。
torchair.scope	npu_stream_switch	图执行过程中，指定图内多个算子分发到不同stream做并行计算。
	npu_wait_tensor	图执行过程中，控制图内多stream并行计算时序。
	super_kernel	图执行过程中，标记图内能融合为SuperKernel的上下文算子范围。
	limit_core_num	图执行过程中，指定图内范围内的算子执行时的AI Core数和Vector Core数。
torchair.llm_datadist	create_npu_tensors	通过一串Device地址创建PyTorch在NPU上的Tensors。主要使用场景是用于创建大模型中的KV Cache Tensors，所有KV Cache的shape和dtype都一致。

父主题： API参考