PyTorch Internals: Tensor Storage, Views, and the Dispatch Mechanism

Summary: Why is x.view(...) nearly free? How does torch.matmul know to call the CUDA kernel vs the CPU one? The answers live in three core concepts: the Tensor data model (storage/view/stride), the Dispatch mechanism, and the ATen operator library underneath. This post makes them clear with reproducible experiments—understanding these makes debugging memory leaks, performance bottlenecks, and custom ops an order of magnitude easier.

0. Experiment Goals

• Dissect the three-layer Tensor data model: storage (data) → view (view) → stride (stride)
• Verify "view is O(1)": changing metadata does not copy data
• Understand the Dispatch mechanism: how one op name routes to CPU/CUDA/custom backends
• Derive engineering practice from internals: memory layout, operator fusion, custom kernels

1. The Tensor Data Model: storage / view / stride

import torch

# A 4x4 tensor: data lives in storage; shape/strides are "view metadata"
x = torch.arange(16, dtype=torch.float32).reshape(4, 4)
print("storage:", x.untyped_storage().data_ptr())   # the single physical home of data
print("shape:", x.shape, "| strides:", x.stride())  # (4,1): row stride 4, col stride 1
print("numel (storage elements):", x.untyped_storage().nbytes() // 4)

Step-by-step:

• Storage: the single physical location of tensor data (a contiguous memory block). data_ptr() is its address in memory.
• Shape + Stride: the "view" that interprets the storage. stride=(4,1) means "jump 4 elements per row, 1 per column"—the addressing formula x[i][j] = storage[i*4 + j].
• Key insight: reshape/view usually only change shape/strides metadata, not data—that is why view ops are "nearly free".

2. Experiment: view Is O(1) (Transpose Does Not Copy)

# Transpose: reverse strides, data untouched
xt = x.t()                      # 4x4 → 4x4 (a view)
print("t strides:", xt.stride())                # (1,4) — only metadata changed
print("same storage:", xt.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())

# Contiguity: non-contiguous views force many ops to .contiguous() first (real copy)
print("x contiguous:", x.is_contiguous(), "| xt contiguous:", xt.is_contiguous())

# .contiguous() is where the real copy happens
xtc = xt.contiguous()
print("different storage after copy:", xtc.untyped_storage().data_ptr() != x.untyped_storage().data_ptr())

Results:

t strides: (1, 4)
same storage: True
x contiguous: True | xt contiguous: False
different storage after copy: True

How it works:

• t() just swaps strides from (4,1) to (1,4)—an O(1) metadata operation, no data touched.
• But non-contiguous tensors cannot be used directly by many ops (some BLAS, non-contiguous view); they must .contiguous()—this is where the real copy happens.
• Engineering meaning: transpose followed by a matmul makes PyTorch implicitly copy once. Avoid it by planning strides from the start (permute + the right memory layout like channels_last)—a real win when optimizing CNN/Transformer performance.

3. The Dispatch Mechanism: How One Op Picks a Kernel

# The same torch.add: CPU tensors take the CPU kernel, CUDA tensors take the CUDA kernel
import torch

cpu_x = torch.randn(3, 3)
cuda_x = cpu_x.cuda()

# Dispatch keys: dtype + device + requires_grad
# torch.add → ATen add() → dispatcher routes by key
#   CPU tensor     → add_cpu (CPU impl)
#   CUDA tensor   → add_cuda (CUDA impl)
#   requires_grad=True → hits the autograd key first, then forwards to the backend impl

How it works:

• The real implementation of PyTorch ops lives in the ATen operator library (C++). torch.add at the Python layer is just an entry point; routing happens in the Dispatcher.
• The Dispatcher uses a set of keys to pick the implementation. Common keys: DispatchKey::CPU, DispatchKey::CUDA, DispatchKey::Autograd, DispatchKey::Meta (shape-only, no execution).
• Keys are layered: with requires_grad=True, the Autograd key fires first (builds the backward graph), then drilling down to the CPU/CUDA compute key.
• This explains why custom ops/custom backends are easy to integrate: you only register an implementation for your key (e.g., register add for a custom device), the dispatcher routes automatically, and upper-layer code changes nothing.

4. Experiment: Register a Custom Op via torch.library (Through Dispatch)

from torch.library import Library, impl

# Register a custom op on Meta and CPU keys
m = Library("mylib", "DEF")           # DEF: define a new op
m.define("triple(Tensor x) -> Tensor")  # declare the op signature

@impl(m, "triple", "CPU")
def triple_cpu(x):
    return x * 3

@impl(m, "triple", "Meta")            # meta: shape-only (for compilation/shape inference)
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.]) — routed to the CPU key
print(torch.ops.mylib.triple(torch.empty(2, 2, device="meta")).shape)  # meta key

Results:

tensor([3., 6., 9.])
torch.Size([2, 2])

Engineering notes:

• Custom ops = registering dispatch implementations: register on CPU/CUDA/Meta keys and the op is recognized by torch.compile and composes with autograd.
• Kernel optimization (Triton/CUDA) plugs into the PyTorch ecosystem through this entry point—no upper-layer training code changes.
• Layered dispatch also structures debugging: performance issues → backend (CPU/CUDA) implementations; gradient issues → the Autograd key's graph construction.

5. Engineering Practice Derived from Internals

Memory layout:

• Contiguous vs non-contiguous drives hidden copies in every op. torch.compile/compiler optimizations (see KB Compile Stack) fuse layout conversions automatically, but layout planning is your responsibility when writing kernels by hand.
• Memory-leak triage: first look for hidden .contiguous() copies (in torch.profiler they appear as heavy copy_).

Operator fusion:

• view/reshape/t are free—compilers can fuse them into neighboring ops without penalty. This is the confidence behind torch.compile's bold graph rewriting.
• Custom kernels should read/write once, compute many times (tiling), minimizing intermediate storage allocation.

Dispatch & custom backends:

• Adding new hardware (Ascend/Cambricon): just implement the op on your custom DispatchKey; upper-layer PyTorch code stays compatible.
• Quantized/sparse kernels also register through dispatch, decoupled from the framework.

6. Engineering Takeaways

1. Tensor = storage + shape/strides: views are metadata ops, O(1), zero copy.
2. Non-contiguity is a hidden copy source: .contiguous() is the real copy; plan layout from the start.
3. Dispatch is the op router: CPU/CUDA/Autograd/Meta keys each do their job.
4. requires_grad hits the Autograd key first: build the graph, then compute—this is autograd's entry.
5. Custom ops = registering dispatch implementations: a few lines via torch.library, upper layer unchanged.
6. Layered debugging: performance → backend keys; gradients → the autograd key.
7. Hand-written kernels own their layout: the compiler can help, but storage allocation for custom kernels is engineering responsibility.

7. Further Reading

• Inside PyTorch 2 (dispatcher design, official blog)
• torch.library API docs
• KB: 09. Deep Learning Frameworks · Autograd Deep Dive · Compile Stack