Beyond Size and Stride: Unleash Performance with Device-Aware Tensor Layouts

In this talk, we review the many ways, both obvious and less obvious, that tensor layouts affect performance — spanning cache locality, memory bandwidth utilization, NUMA effects, and hardware-specific tiling constraints. PyTorch tensors encode only minimal information about these physical considerations — size, stride, and storage offset — which is insufficient to capture the hardware layouts required by modern accelerators.

To address this, we introduce a tensor layout extension that enables device-aware specialization of tensor layouts, including tiling and NUMA-aware placement. This lightweight extension preserves standard PyTorch tensor semantics by introducing a separate, more expressive physical layout, allowing precise control over these layouts without modifying existing PyTorch code.

Finally, we report our experiences using this approach in practice. We demonstrate how torch.compile (Inductor) leverages these controls to automatically adapt layouts to the target device and transparently adjust computations. We discuss the resulting performance benefits, limitations, and lessons learned.