PyTorch profiler custom backend Ops

在pytorch这个最流行的AI框架上注册新的后端已经成为了部分芯片推理模型的重要方式，可以极大的减少研发芯片专属的推理引擎的花费，只需要注册一个新的后端，然后在算子库中实现自己的硬件强相关的算子实现，就可以用最少的精力开始算子正确性等的验证，而不需要等完整的芯片编译器完工。

一切的前提

你需要知道在注册算子的时候，你的算子是绑在哪个Key上(PrivateUse1,PrivateUse2,PrivateUse3..)，类似下面的一些注册

1	torch.utils.rename_privateuse1_backend("xxnpu")

1
2
3

TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
    m.impl("add.Tensor", my_add_kernel);
}

torch.profiler工具

首先来看看这个工具能帮我们做什么事情

CPU 与 GPU 算子耗时：精确记录每个 PyTorch 算子（如 aten::linear, aten::bmm）在 CPU 端的 Launch 耗时，以及在 GPU 端的实际执行耗时。
CUDA Kernel 映射：将 PyTorch 的高级算子映射到底层具体的 CUDA Kernel（如 flash_fwd_kernel），让你知道到底在跑什么底层代码。
显存分配与释放：记录显存的申请和释放事件，追踪是哪个算子导致了显存峰值（OOM 排查神器）。
分布式通信 (NCCL)：在分布式训练/推理中，抓取 AllReduce、AllGather 等通信算子的耗时。
CPU-GPU 同步与 Gap：通过对比 CPU 线程和 GPU Stream 的时间线，找出 CPU 是否“喂不饱” GPU（即 GPU 在等 CPU 发指令的空闲时间）。

当然，我们的自定义算子也同样可以使用这个工具来完成trace。

简单实践

def trace_handler(p):
	p.export_chrome_trace("./logs/paged_attention_llama3_8b_decode_shape.json")

with profile(
	activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],  # xxnpu
	schedule=torch.profiler.schedule(wait=1, warmup=1, active=9,repeat=1),
	record_shapes=True,
	with_stack=True,
	profile_memory=True,
	on_trace_ready=trace_handler,
) as prof:
	for _ in range(10):
		output = torch.ops.xxnpu.add(a, b)
		prof.step()
print(prof.key_averages(group_by_input_shape=True).table(
	sort_by="self_cpu_time_total", row_limit=80
))

trace_handler : 一个记录trace结果的回调函数，一定要写，这样我们可以后续在网站上可视化我们的trace结果
schedule：对于多轮迭代的任务，这个配置可以来指定几次是静默执行(wait)，几次是warmup，后面要记录多少次的trace结果，迭代次数充足的情况下重复多少次前面的三种行为。
prof.step()：多轮迭代中需要明确告知profiler，你已经执行完了一次，配合schedule使用

按照上述代码中的打印，出来的结果大概是这个样子

------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------  
              Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls                                                          Input Shapes  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------  
         aten::add        98.64%        7.330s        98.64%        7.330s     814.396ms           0 B           0 B             9      [[2560, 16, 8, 128], [2560, 16, 8, 128], [], [2560, 16, 8, 128]]  
    aten::_to_copy         1.29%      96.164ms         1.31%      96.979ms       5.388ms           0 B           0 B            18                          [[2560, 16, 8, 128], [], [], [], [], [], []]  
     ProfilerStep*         0.04%       2.938ms       100.00%        7.430s     825.608ms           0 B           0 B             9                                                                    []  
       aten::empty         0.02%       1.237ms         0.02%       1.237ms      45.806us           0 B           0 B            27                                              [[], [], [], [], [], []]  
         aten::add         0.01%     400.446us        98.65%        7.330s     814.488ms           0 B           0 B             9                          [[2560, 16, 8, 128], [2560, 16, 8, 128], []]  
          aten::to         0.00%     161.684us         1.31%      97.140ms       5.397ms           0 B           0 B            18                      [[2560, 16, 8, 128], [], [], [], [], [], [], []]  
------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  --------------------------------------------------------------------  
Self CPU time total: 7.430s

结果中可以看到存在两个aten::add，但是他们的input shape其实是不同的，因为add算子被dispatch到了 aten::add.out上了，最后一个shape就是表示output shape的。
需要注意一点，ProfilerActivity.CPU是一定要加的，因为算子都是从CPU上逐个dispatch到不同后端的，如果不加的话，是获取不到任何有用的信息的。

PyTorch profiler custom backend Ops

一切的前提

torch.profiler工具

简单实践

进阶使用 [TODO]