[feat] npu support enable_torch_compile

[XDaoHong]

Motivation

TorchAir (Torch Ascend Intermediate Representation) is an extension library that provides graph mode capabilities for torch_npu. It enables users to perform graph-mode inference on NPU using PyTorch and torch_npu. TorchAir externally offers a torch.compile backend for NPU, which interfaces with torch._dynamo. Through the following features, performance optimization and capability enhancement of the torch fx graph can be achieved.

Main Features:

1. Basic Features:

• Enable NPU kernels that depend on host-value tiling operators (e.g., FIA) to support npugraph
• Graph input copy optimization
• Memory reuse across multi-graphs

2. FX Pass:

• In-place optimization
• Redundant operator elimination
• NPU fused operator passes

3. Advanced Features:

• Static shape kernel compilation
• Multi-stream within single graphs
• Compilation caching

Modifications

1. Rewrite the capture function;
2. Encapsulate the kvcache input (input needs all kvcache);
3. Pad the block table to the max length;
4. TorchAir input preparation;

The calling process is as follows.

Accuracy Tests

python3 few_shot_gsm8k.py --data-path "/path/to/model/test.jsonl.txt” --parallel 32 --num-questions 200

Accuracy: 0.865
Invalid: 0.000
Latency: 43.077 s
Output throughput: 795.877 token/s

Benchmarking and Profiling

Future roadmaps

In the torch_npu 7.2.0 version, the reduce-overhead mode of the torchair backend will support torch.compile(model, dynamic=True). This mode will be set as the default in get_compile_backend(), enabling support for methods wrapped by the @torch.compile() decorator.
In the torch_npu 7.3.0 version, the capture and replay of NPUGraph currently integrated in the torchair backend will be changed to optional execution. The torchair backend will only perform optimizations such as fx pass optimization and static kernel compilation, while the capture and replay of NPUGraph will be implemented independently. This design is closer to the implementation of CudaGraphRunner, decoupling fx graph optimization from graph offloading.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ssshinigami]

Could you please add description and motivation
And accuracy and performance measurements

[iforgetmyname]

please add a comment "TODO(iforgetmyname): Change this default value once CANN version 8.3.RC1" to help me remember to change this default value to reduce-overhead

[XDaoHong]

done

[ZhengdQin]

revert this part

[eshoguli]

General comment about NPUGraphRunner updates. NPUGraphRunner is inherited from CudaGraphRunner and reuse capturing (ForwardBatch instantiation, initialization and capturing functionality) and partially replay pipeline from inherited type. NPUGraph python type is used.

Your changes don't use anything from that. You have custom implementation of capturing and don't use NPUGraph python type for inference. As result, you need to implement separate runner: TorchAirRunner or something like that.

Can you, please, explain offline/online: why you need to use NpuGraphRunner? thanks!