PyTorch 内部架构：Tensor 存储、View 与 Dispatch 机制
摘要:
x.view(...)为什么几乎不花时间？torch.matmul怎么知道该调 CUDA 还是 CPU 内核？答案都在 PyTorch 的三个内核概念里：Tensor 数据模型（storage/view/stride）、Dispatch 分发机制、以及背后的 ATen 算子库。本文用可复现实验把它们讲透——理解了这些，调试显存泄漏、性能瓶颈和自定义算子会轻松一个量级。
0. 实验目标
- 拆开 Tensor 的三层数据模型：storage（数据）→ view（视图）→ stride（步长）
- 验证"view 是 O(1) 操作"：改元数据不拷贝数据
- 理解 Dispatch 机制：一个算子名如何路由到 CPU/CUDA/自定义后端
- 从内部机制反推工程实践：内存布局、算子融合、自定义内核
1. Tensor 数据模型：storage / view / stride
python
import torch
# 一个 4x4 张量：数据存在 storage 里，shape/strides 是"视图元数据"
x = torch.arange(16, dtype=torch.float32).reshape(4, 4)
print("storage:", x.untyped_storage().data_ptr()) # 数据唯一存放处
print("shape:", x.shape, "| strides:", x.stride()) # (4,1)：行步长 4，列步长 1
print("numel(存储元素):", x.untyped_storage().nbytes() // 4)逐段拆解：
- Storage：张量数据的唯一物理存放处（连续内存块）。
data_ptr()是它在内存里的地址。 - Shape + Stride：解释 storage 的"视图"。
stride=(4,1)表示"行方向跳 4 个元素、列方向跳 1 个元素"——这就是x[i][j] = storage[i*4 + j]的寻址公式。 - 关键洞察：
reshape/view多数时候只改 shape/strides 元数据，不搬数据——这就是 view 操作"几乎零成本"的原因。
2. 实验：view 是 O(1) 的（转置不拷贝）
python
# 转置：strides 反转，数据原地不动
xt = x.t() # 4x4 → 4x4（视图）
print("t strides:", xt.stride()) # (1,4) —— 只换了元数据
print("同一 storage:", xt.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())
# 连续性：视图非连续时，很多算子要先 .contiguous()（会触发拷贝）
print("x 连续:", x.is_contiguous(), "| xt 连续:", xt.is_contiguous())
# .contiguous() 才真正拷贝
xtc = xt.contiguous()
print("拷贝后不同 storage:", xtc.untyped_storage().data_ptr() != x.untyped_storage().data_ptr())实验结果：
t strides: (1, 4)
同一 storage: True
x 连续: True | xt 连续: False
拷贝后不同 storage: True原理讲解：
t()只是把 strides 从(4,1)换成(1,4)——O(1) 元数据操作，不碰数据。- 但非连续张量在很多算子（如部分 BLAS、
view非连续场景）下不能直接用，必须.contiguous()——这一步才发生真实拷贝。 - 工程含义：
transpose后立刻做矩阵乘，PyTorch 会在背后隐式拷贝一次。想避免，用permute+ 正确的内存布局（channels_last等）从一开始就规划好 strides——这在优化 CNN/Transformer 性能时是真实收益点。
3. Dispatch 机制：一个算子如何选内核
python
# 同一个 torch.add，CPU 张量走 CPU 内核，CUDA 张量走 CUDA 内核
import torch
cpu_x = torch.randn(3, 3)
cuda_x = cpu_x.cuda()
# dispatch 依据：dtype + device + 是否 requires_grad
# torch.add → ATen 的 add() → dispatcher 按 key 路由
# CPU 张量 → add_cpu（CPU 实现）
# CUDA 张量 → add_cuda（CUDA 实现）
# requires_grad=True → 先进入 autograd key，再转发到后端实现原理讲解：
- PyTorch 算子的真身在 ATen 算子库（C++）。Python 层的
torch.add只是入口，真正的路由靠 Dispatcher。 - Dispatcher 用 key 集合 决定调哪个实现，常见 key：
DispatchKey::CPU、DispatchKey::CUDA、DispatchKey::Autograd、DispatchKey::Meta（仅算 shape，不执行）等。 - key 是分层处理的：
requires_grad=True时先命中Autogradkey（构建反向图），再下钻到CPU/CUDA的实际计算 key。 - 这解释了为什么"自定义算子/自定义后端"容易接入：你只需要为某个 key 注册实现（比如为自定义设备注册
add），dispatcher 自动路由，上层代码零改动。
4. 实验：用 torch.library 注册自定义算子（走 dispatch）
python
from torch.library import Library, impl
# 注册一个自定义算子在 meta/CPU 两个 key 上的实现
m = Library("mylib", "DEF") # DEF：定义新算子
m.define("triple(Tensor x) -> Tensor") # 声明算子签名
@impl(m, "triple", "CPU")
def triple_cpu(x):
return x * 3
@impl(m, "triple", "Meta") # meta：只算 shape 不执行（用于编译/形状推导）
def triple_meta(x):
return x.new_empty(x.shape)
x = torch.tensor([1.0, 2.0, 3.0])
print(torch.ops.mylib.triple(x)) # tensor([3., 6., 9.]) — CPU key 被路由
print(torch.ops.mylib.triple(torch.empty(2, 2, device="meta")).shape) # meta key实验结果：
tensor([3., 6., 9.])
torch.Size([2, 2])工程启示：
- 自定义算子 = 注册 dispatch 实现：注册到
CPU/CUDA/Meta各 key，算子就能被torch.compile识别、被 autograd 组合。 - 内核优化（Triton/CUDA）通过这个入口无缝接入 PyTorch 生态——不用改上层训练代码。
- Dispatch 分层也让调试更有条理：性能问题看
CPU/CUDAkey 的实现，梯度问题看Autogradkey 的图构建。
5. 从内部机制反推的工程实践
内存布局：
- 连续 vs 非连续影响每个算子的隐藏拷贝。
torch.compile/编译器优化（本知识库 编译栈 篇）会自动融合布局转换，但手写内核时布局规划是自己的责任。 - 显存泄漏排查：先看是否有隐藏的
.contiguous()拷贝（torch.profiler里表现为copy_大量出现）。
算子融合：
view/reshape/t零成本——编译器把它们"融合"进相邻算子毫无负担。这也是torch.compile敢大胆重构图的底气。- 自定义内核时尽量一次读写、多次计算（Tile 化），减少中间张量的 storage 分配。
dispatch 与自定义后端：
- 接入新硬件（昇腾/寒武纪等）：只需实现算子在自定义
DispatchKey上的实现，上层 PyTorch 代码兼容。 - 量化/稀疏内核同样走 dispatch 注册，与框架解耦。
6. 工程启示汇总
- Tensor = storage + shape/strides：view 是元数据操作，O(1) 零拷贝。
- 非连续是隐藏拷贝源：
.contiguous()才是真实拷贝；布局从一开始就规划。 - Dispatch 是算子的路由器：CPU/CUDA/Autograd/Meta 各 key 各司其职。
- requires_grad 先走 Autograd key：先建图、后计算，这就是 autograd 的入口。
- 自定义算子 = 注册 dispatch 实现：
torch.library几行接入，上层零改动。 - 调试分层：性能看后端 key、梯度看 autograd key。
- 手写内核要自己管布局：编译器能帮你，但自定义 kernel 的 storage 分配是工程责任。