Tensor Layouts and Memory

1. What is it?#

Layout is how a logical tensor’s elements map to physical memory addresses. The same (B, h, S, D) tensor can be stored in 24 different orders, and the choice changes kernel performance by 2-5x.

Logical:  attention output, 4 dimensions
Physical: one flat array of bytes

Layout = the permutation and stride pattern connecting them

2. Why does it exist as a topic?#

Because GPUs read memory in 32-byte or 128-byte transactions serving a whole warp, and whether your 32 threads’ addresses fall in one transaction or 32 determines whether you get 3,350 GB/s or 100 GB/s.

Layout is the difference between a kernel that runs at bandwidth and one that doesn’t. It’s also the reason KV cache design is a real engineering decision rather than an implementation detail.

3. Simple analogy#

A warehouse’s shelving scheme. The inventory is the same either way, but if items that are picked together sit together, a picker fills an order in one pass down one aisle. If they’re scattered, the picker walks the whole warehouse per order.

Memory transactions are the picker’s trips. Layout determines how many trips per order.

4. Tiny example#

import torch, time

B, h, S, D = 8, 32, 4096, 128

def bench(t, label):
    torch.cuda.synchronize(); t0 = time.perf_counter()
    for _ in range(50): s = t.sum()
    torch.cuda.synchronize()
    dt = (time.perf_counter()-t0)/50
    gb = t.numel()*t.element_size()/1e9
    print(f"{label:28s} {dt*1e3:7.2f} ms  {gb/dt:7.0f} GB/s  contiguous={t.is_contiguous()}")

x = torch.randn(B, h, S, D, device='cuda', dtype=torch.float16)
bench(x, "BhSD contiguous")
bench(x.transpose(1, 2), "BShD view (non-contig)")
bench(x.transpose(1, 2).contiguous(), "BShD materialized")

Typical:

BhSD contiguous              1.42 ms     756 GB/s  contiguous=True
BShD view (non-contig)       3.87 ms     277 GB/s  contiguous=False
BShD materialized            1.44 ms     747 GB/s  contiguous=True

Same data, same operation, 2.7x difference purely from access pattern. And the .contiguous() call that fixes it costs a full read+write of the tensor.

5. Technical explanation#

Coalescing — the governing rule#

A warp of 32 threads issues 32 addresses. The memory system services them in as few 32-byte sectors as possible.

PERFECT (coalesced):
  thread 0 → byte 0-3, thread 1 → byte 4-7, ..., thread 31 → byte 124-127
  → 128 contiguous bytes → 4 sectors → 1 memory transaction burst

WORST (strided by 4096 bytes):
  thread 0 → byte 0, thread 1 → byte 4096, ...
  → 32 separate 32-byte sectors → 32 transactions
  → you fetched 1024 bytes to use 128. 8x waste (or 32x for 4-byte elements)

The rule: consecutive threads should read consecutive addresses. Which means: the fastest- varying dimension of your tensor should be the one that consecutive threads index.

KV cache layouts — a real design decision#

OPTION A: [num_blocks, block_size, n_kv_heads, head_dim]
  Appending token t: write n_kv_heads × head_dim contiguous elements. ONE write. ✓
  Attention reading head j across the block: stride of n_kv_heads×head_dim. Strided. ✗

OPTION B: [num_blocks, n_kv_heads, block_size, head_dim]
  Appending token t: write head_dim elements at n_kv_heads different offsets. Scattered. ✗
  Attention reading head j across the block: contiguous. ✓

OPTION C (vLLM's key cache): [num_blocks, n_kv_heads, head_dim/x, block_size, x]
  where x = 16 bytes / element_size (8 for FP16)
  Designed so the attention kernel's per-thread reads are 16-byte vectorized
  and coalesced across the warp.

vLLM uses different layouts for K and V because the attention kernel accesses them differently: K is used in q·kᵀ (reduction over head_dim), V in p·v (reduction over sequence). The layouts are chosen per-usage, not for elegance. Reading vllm/attention/ops/paged_attn.py and the corresponding CUDA is the best way to internalize this.

Weight layouts#

Row-major (out, in):     nn.Linear default. Computing x @ W.T reads W rows contiguously.
Column-major:            what BLAS traditionally wants.
Interleaved/packed:      for quantized weights, 4-bit values packed two per byte,
                         often in an order that makes dequantization vectorizable.
Tensor-core native:      some kernels want weights pre-swizzled to match the MMA
                         fragment layout, avoiding a shuffle in the inner loop.

Marlin (a fast W4A16 kernel) pre-permutes weights at load time into an order that makes the dequantize-and-feed-tensor-core path branch-free and coalesced. The permutation is done once at startup; the payoff is every decode step. This is a general pattern: pay layout cost once at load, not per inference.

Padding and alignment#

Tensor core requirements:
  FP16: dimensions should be multiples of 8 (ideally 64/128 for tiles)
  INT8/FP8: multiples of 16
  Pointers: 16-byte aligned for vectorized loads

Shared memory bank conflicts:
  32 banks of 4 bytes. Threads hitting the same bank serialize.
  Fix: pad the leading dimension of shared-memory tiles by 1 element,
       or use a swizzled (XOR-based) layout.

The classic shared-memory padding trick:

__shared__ float tile[32][33];   // 33, not 32 — breaks the bank-conflict pattern

6. Under the hood#

Measure coalescing directly:

ncu --metrics l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\
l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum ./your_program

The ratio sectors / requests tells you how many 32-byte sectors each warp request touched:

ratio = 4    → perfectly coalesced (128 bytes = 4 sectors per warp)
ratio = 32   → fully uncoalesced (each thread its own sector)

Also useful: l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum for shared-memory bank conflicts.

7. Performance implications#

8. Production implications#

• Convert layouts once at load time, never per request. Weight permutation, KV layout, channels-last conversion — all startup work.
• Match the KV layout to your attention kernel. If you switch attention backends, the layout may need to change.
• Watch for .contiguous() in profiles. Each one is a full tensor copy. Some are necessary; many are accidental.
• Pad vocabulary and hidden dimensions to tensor-core-friendly multiples. Models are usually designed this way; custom heads often aren’t.
• When adding a custom kernel, measure the sectors/requests ratio before optimizing anything else.

9. Common mistakes#

Choosing a layout for readability. The kernel doesn’t care about your aesthetics.

Forgetting that transpose is free but the next op may not be. The cost is deferred, not avoided.

Assuming PyTorch will pick a good layout. It preserves what you give it.

Copying a KV layout from a paper without checking your kernel’s access pattern.

Bank conflicts in hand-written shared-memory tiles. Pad or swizzle.

Ignoring alignment for vectorized loads. A half8 load requires 16-byte alignment; without it you silently get scalar loads.

10. Hands-on exercise#

A. Measure coalescing. Write two CUDA (or Triton) kernels that sum a large 2D array: one reading row-wise, one column-wise. Measure achieved bandwidth and the sectors/requests ratio for each with ncu.

B. KV layout study. Implement a simple decode-attention kernel for layouts A and B from section 5. Benchmark both at S=4096, h_kv=8, D=128. Which wins, and does it match your prediction?

C. Bank conflicts. Write a shared-memory transpose with tile[32][32] and tile[32][33]. Measure both and check the bank-conflict counter.

D. Hunt for .contiguous(). Profile a model forward pass and find every copy_ kernel. Trace each back to its cause in the Python code. How many are avoidable?

E. Alignment. Allocate a tensor with a deliberately misaligned offset (slice by 1 element) and benchmark a vectorized kernel on it vs the aligned version.

11. Interview questions#

1. What is memory coalescing and what is the worst-case penalty?
2. Design a KV cache layout. Justify every dimension ordering.
3. Why does vLLM use different layouts for K and V?
4. What is a shared-memory bank conflict and how do you avoid it?
5. Why do quantized-GEMM kernels pre-permute their weights?
6. When is transpose() free and when does it cost you?
7. How would you measure whether a kernel’s memory accesses are coalesced?

12. Further reading#

• [REFERENCE] CUDA C++ Best Practices Guide, “Coalesced Access to Global Memory”
• [REFERENCE] vLLM csrc/attention/attention_kernels.cu — read the layout comments
• [ESTABLISHED] Marlin kernel source and design notes
• [REFERENCE] Nsight Compute memory workload analysis section
• Next: 08 — Kernels and kernel launches