Pytorch知识记录

前言/引言

随着LLM的快速发展，对于各个硬件厂商来说快速的迭代验证某些功能是非常必要的，比如算子。如何能够快速的验证自家算子实现的是否准确，最简单的办法就是在torch中增加自定义device，然后利用torch的dispatch机制来逐个调用，这样不需要等待编译器完工就可以跑起来整个E2E的算子精度验证，当然这只是一种较为基础的验证。

func schema

Pytorch中支持的算子很多很多，而且每个算子都存在多种使用方式，可以参考pytorch/aten/src/ATen/native/native_functions.yaml中的函数定义，下面以fmod为例来说明集中不同的schema。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

- func: fmod.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  dispatch:
    CompositeExplicitAutograd: fmod_out
  tags: pointwise

- func: fmod.Scalar(Tensor self, Scalar other) -> Tensor
  device_check: NoCheck   # TensorIterator
  variants: method, function
  dispatch:
    CompositeExplicitAutograd: fmod
  tags: [core, pointwise]

- func: fmod_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  variants: method
  dispatch:
    CompositeExplicitAutograd: fmod_
  tags: pointwise

- func: fmod.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  structured: True
  structured_inherits: TensorIteratorBase
  dispatch:
    CPU, CUDA, MPS, MTIA, XPU: fmod_out
  tags: pointwise

- func: fmod.Tensor(Tensor self, Tensor other) -> Tensor
  device_check: NoCheck   # TensorIterator
  structured_delegate: fmod.Tensor_out
  variants: method, function
  tags: [core, pointwise]

- func: fmod_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  variants: method
  structured_delegate: fmod.Tensor_out
  tags: pointwise

这段 YAML 是 PyTorch 的算子注册文件 native_functions.yaml 中的几个条目，定义了 fmod（浮点取模）算子的各种变体。我们逐项拆开看，然后再总结它们和自定义后端的关系。

1. 算子签名 (func:)

1

- func: fmod.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)

• fmod.Scalar_out：算子名字，fmod 是基础名，.Scalar_out 是重载名。表示这是 第二个操作数是标量，且结果写入预分配张量 out 的版本。
• 参数列表：
  • Tensor self：输入张量。
  • Scalar other：一个标量值。
  • * 后面的都是仅关键字参数（keyword-only），即 Tensor(a!) out。
  • Tensor(a!)：a! 表示这个张量会被原地修改，且函数结束后返回同一个张量。
• 返回值：-> Tensor(a!)，即返回被修改后的 out 张量。

其他变体同理：

• fmod.Scalar：标量版本，返回新张量（没有 out，也不原地修改）。
• fmod_.Scalar：标量版本，原地修改 self（Tensor(a!) self）。
• fmod.Tensor_out：两个张量相除取模，结果写入 out。
• fmod.Tensor：两个张量，返回新张量。
• fmod_.Tensor：两个张量，原地修改 self。

2. device_check: NoCheck

1

device_check: NoCheck   # TensorIterator

• 默认情况下，PyTorch 会检查一个算子的所有输入张量是否在同一个设备上（比如不能一个在 CPU，一个在 CUDA）。
• NoCheck 表示不做这个检查。因为 fmod 内部使用了 TensorIterator，它自己会处理设备不匹配的情况（例如会报出更合适的错误信息），所以这里提前关闭了框架层的设备一致性检查。
• 注释里的 # TensorIterator 也说明了它依赖这个机制。

3. dispatch:（分发规则）

这是最关键的部分，决定了不同 dispatch key 下调用哪个 C++ 函数。

① 标量版本的 dispatch

1
2

dispatch:
  CompositeExplicitAutograd: fmod_out

• 标量版的 out、普通、in-place 三个变体，其 dispatch 键都只写了 CompositeExplicitAutograd。
• 这表示：这些算子被归类为“组合显式自动微分”算子。它们的前向实现（fmod_out、fmod、fmod_）是由其他基础算子组合而成的，而不是直接手写的底层 kernel。
• 对 autograd 而言，CompositeExplicitAutograd 意味着该算子有自己显式的反向传播规则（在 derivatives.yaml 中定义），但它的前向并不直接调用设备特化 kernel，而是调用组合实现（比如把标量转成张量再调张量版本，或直接使用循环）。
• 对于你的自定义后端：你不需要为这些标量版本单独注册 kernel。因为 dispatcher 找不到你的后端专属实现时，就会回退到 CompositeExplicitAutograd 的实现，这个实现会转而调用更底层的张量版本或其他算子，最终走到你的后端基础算子（比如 fmod.Tensor_out）。

② 张量版本的 dispatch

1
2

dispatch:
  CPU, CUDA, MPS, MTIA, XPU: fmod_out

• 这里明确列出了 CPU、CUDA、MPS、MTIA、XPU 这几个后端，它们都会直接使用同一个内核 fmod_out。
• 这个 fmod_out 内核很可能是基于 TensorIterator 的通用、高性能实现，因此多个后端可以共享同一套代码（通过 TensorIterator 自动适配设备）。
• 对于你的自定义后端：因为你使用的是 PrivateUse1，而这个列表里没有 PrivateUse1，所以 PyTorch 的 dispatcher 在调用 fmod.Tensor_out 时，找不到你设备专用的 kernel。会发生什么？
  • 如果没有设置任何 fallback，会报错：“no kernel registered for PrivateUse1”。
  • 如果你在注册时使用了 torch._C._dispatch_fallback(torch._C.DispatchKey.PrivateUse1, torch._C.DispatchKey.CompositeImplicitAutograd) 之类的 fallback 策略，则可能会回退到 CPU 实现（导致性能和设备问题），或者回退到组合实现。
  • 推荐做法：如果你的后端可以复用 TensorIterator 的现有 CPU/CUDA 实现，你应该在注册时把 fmod.Tensor_out 的 kernel 也“指向”同一个 fmod_out 函数，即：

    1	m.impl("fmod.Tensor_out", at::native::fmod_out);

    这样你的设备就能直接使用通用的 TensorIterator 内核了。

4. variants:（生成哪些 Python 接口）

1

variants: method, function

• method：为张量生成一个方法，如 tensor.fmod(other)。
• function：生成一个 torch 模块下的函数，如 torch.fmod(tensor, other)。
• 如果没有 variants 字段，默认同时生成 method 和 function（有些版本默认不同，这里显式写了两种）。

5. tags:（标签）

1

tags: [core, pointwise]

• core：标记为核心算子，会进入 PyTorch 的“core ATen IR”中，一些分析工具和编译器（如 torch.compile）会重点关注这些算子。
• pointwise：逐元素操作，意味着输出和输入形状相同，且每个元素独立计算，不依赖相邻元素。这个标签允许图优化器对这类算子进行融合等优化。
• tags: pointwise 单独出现时，表示非核心的逐元素操作。

6. 结构化内核相关 (structured, structured_inherits, structured_delegate)

1
2

structured: True
structured_inherits: TensorIteratorBase

• 这是 PyTorch 的**“结构化内核”机制**，用于减少样板代码。
• structured: True 表示这个算子是一个结构化算子，其内核实现遵循特定的函数签名：void fmod_out_impl(...)，并且代码生成器会自动创建 fmod_out、fmod.Tensor、fmod_.Tensor 等变体，以及自动处理反向传播的注册。
• structured_inherits: TensorIteratorBase 表示这个结构化内核继承自 TensorIteratorBase，因此可以直接使用 TensorIterator 来并行化、向量化、设备适配等。
• structured_delegate: fmod.Tensor_out：对于 fmod.Tensor 和 fmod_.Tensor 来说，它只是委托给 fmod.Tensor_out 的实现。也就是说，fmod.Tensor 的代码生成会直接调用 fmod.Tensor_out 的内核，省去再写一遍的麻烦。

7. 各变体的关系总结

对你的自定义后端意味着什么？

1. 标量版本你不需要管。它们被标记为 CompositeExplicitAutograd，PyTorch 会确保它们通过组合方式工作（最终会调用张量版本）。
2. 张量版本 (fmod.Tensor_out) 是你需要关注的。如果它的 dispatch 列表里没有 PrivateUse1，你需要显式注册一个内核，才能让你的设备支持 torch.fmod 以及 tensor.fmod_() 等。
3. 由于 fmod.Tensor_out 的现有内核可能已经高度泛化（基于 TensorIterator），你可以直接复用那个通用实现，像这样：

   1 2 3

   TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { m.impl("fmod.Tensor_out", &at::native::fmod_out); // 借用 ATen 的通用实现 }

   前提是你的设备支持 TensorIterator（即你实现了分配器、拷贝、基础循环等）。
4. 如果你没有为 PrivateUse1 注册 fmod.Tensor_out，也没有设置合适的 fallback，那么调用 torch.fmod 时会直接报错。

简而言之，这段 YAML 告诉你：张量版本的 fmod 是有具体后端实现的，你必须为你的后端注册它（或利用 fallback 绕过），而标量版本会自动组合出来，不用你动手。