Skip to content

Memory Formats & Layouts

Overview

Memory format (contiguous vs channels_last vs blocksparse) determines how kernels read/write tensors in DRAM. The same logical tensor can be 2-4x faster to process— or to pass between layers— based purely on physical layout. GPUs love coalesced, aligned, large-block accesses.

  • NCHW (default): channels innermost → channel-heavy ops (matmul, batchnorm) walk memory efficiently.
  • NHWC / channels_last: spatial dims innermost → convolution/attention gather favorably; torch.compile+channels_last is the modern fast path.
  • memory_format is metadata on a tensor: .to(memory_format=torch.channels_last) reshapes layout (copy), kernels choose paths by it.
  • Coalescing: neighboring threads read neighboring addresses → one 128-byte transaction.

Modern GPUs: increasing batch/linear dims contiguous is king; permute/transpose in hot paths silently triggers full copies.

Viewing the Format & Why It Matters

import torch

x = torch.randn(8, 3, 64, 64) # NCHW contiguous
print(x.is_contiguous(memory_format=torch.channels_last)) # False

x_cl = x.to(memory_format=torch.channels_last) # copy into NHWC
print(x_cl.is_contiguous(memory_format=torch.channels_last)) # True
print(x_cl.shape) # same logical shape!

The tensor is logically identical; only the physical order of elements changed.

-

NCHW vs NHWC in Layers

import torch, torch.nn as nn, time

def bench_conv(N, C, H, W, fmt, n=20):
 m = nn.Conv2d(C, C, 3, padding=1).cuda()
 m = m.to(memory_format=fmt)
 x = torch.randn(N, C, H, W, device='cuda').to(memory_format=fmt)
 for _ in range(3): m(x)
 torch.cuda.synchronize()
 t0 = time.perf_counter()
 for _ in range(n): m(x)
 torch.cuda.synchronize()
 return (time.perf_counter() - t0) / n * 1e3

cf = torch.contiguous_format
cl = torch.channels_last
print("contiguous: %.2f ms" % bench_conv(16, 64, 64, 64, cf))
print("channels_last: %.2f ms" % bench_conv(16, 64, 64, 64, cl))
# On many GPUs channels_last is faster for convs (memory coalescing).

Results vary by GPU/version— always measure. The direction: convolutions + AMP favor channels_last; pure matmul/attention favors plain contiguous.

When Formats Matter Most

Op Best format Why
Conv2d/3d channels_last often coalesced spatial reads
BatchNorm channels_last w/ conv fused layout preserved end-to-end
Linear/Matmul contiguous GEMM wants row-major
Attention (QKV gather) contiguous (or fused kernels) index arithmetic simple
AMP (fp16) BMMs contiguous alignment

Format consistency across a model

model = model.to(memory_format=torch.channels_last)
# inputs too:
x = x.to(memory_format=torch.channels_last)
# avoid per-layer to() conversions in the hot path

Pitfalls

# 1. to() back and forth silently copies
y = x_cl.to(memory_format=torch.contiguous_format) # full copy

# 2. Weights and activations must match format for fusion
# torch.compile handles conversions automatically; eager doesn't.

# 3. Some ops don't support channels_last -> forced fallback copies
# check via

-

torch.compile Makes Most of This Automatic

torch.compile(model, mode="max-autotune")
# Inductor picks layouts per kernel; you mostly stop hand-tuning formats.

Prefer torch.compile default path over manual channels_last surgery in eager code. Manual formats still matter for: custom kernels, ONNX/export paths, and memory-constrained inference.

Strategy Checklist

  1. Default: contiguous + torch.compile (max-autotune) and let Inductor choose.
  2. Conv-heavy eager nets: channels_last for model AND data.
  3. Never convert format inside a hot loop; convert once at the boundary.
  4. Profile with torch.profiler— a "copy_" node next to a layer signals format mismatch.
  5. Export (Ch 07): pin formats in the export inputs.

-

Key Takeaways

  • Format = physical layout; changing it costs a full copy.
  • channels_last helps conv-heavy and some AMP paths; contiguous is king for GEMMs.
  • torch.compile largely automates layout choice— don't hand-tune before measuring.
  • Keep formats consistent end-to-end (weights + activations + inputs).
  • Profile for copy_/to nodes to detect stray conversions.

-