Tensor Memory Accountant

The reading is a page. This is the reps. Do all three blocks on paper first, then check with PyTorch if you have it. The point is to be right before print(x.nbytes).

Block 1 — The original

x = torch.randn(16, 1024, 4096, dtype=torch.float16, device="cuda")
a = x.transpose(1, 2)
b = a.contiguous()
c = b.reshape(16 * 4096, 1024)
d = c.to(torch.float32)
e = d[:, :512]
f = e.clone()

Deliverable: one row per name.

name	shape	dtype	storage bytes	view or copy	keeps alive

Then:

1. Live tensor storage after f, nothing freed. Binary units.
2. First unavoidable full copy.
3. Copy you could delete by changing the next op.
4. Bytes moved by d = c.to(torch.float32) (read + write).

Block 2 — The "it's just a view" PR

w = torch.randn(4096, 4096, dtype=torch.bfloat16, device="cuda")  # weights
h = torch.randn(8, 4096, dtype=torch.bfloat16, device="cuda")
y = h @ w.t()
z = y.view(2, 4, 4096)
p = z.permute(1, 0, 2)
q = p.reshape(4, 8192)

1. Which of y, z, p, q allocate?
2. w.t() — view or copy? Does the matmul care?
3. If the next kernel needs q contiguous, where does the copy land, and how many bytes?

Block 3 — The silent cast

ids = torch.randint(0, 32000, (4, 2048), device="cuda")  # int64 default
emb = torch.randn(32000, 4096, dtype=torch.bfloat16, device="cuda")
hidden = emb[ids]          # gather
hidden = hidden.float()    # "for stability"
logits = hidden @ torch.randn(4096, 32000, dtype=torch.float32, device="cuda")

1. Size of emb, hidden before the cast, hidden after the cast.
2. Is emb[ids] a view into emb? (Be precise about gather.)
3. The author says the float cast "shouldn't matter, it's the same shape." Write the traffic ratio vs staying in bf16 for that line.

Rules

1. float16/bfloat16 = 2 bytes, float32 = 4, int64 = 8.
2. Transpose/permute/view are views when they only rewrite stride.
3. .contiguous(), .clone(), dtype change, and gather allocate.
4. reshape is a view if the storage is already contiguous in that order; otherwise it copies.

Acceptance

1. Block 1: x is 128 MiB. a is a view. b copies 128 MiB. d copies to 256 MiB. e is a view of d. f copies 128 MiB. Live storage includes x, b/c (same storage), d, f.
2. Block 2: w.t() is a view; y allocates; z is a view of y; p is a view; q may copy because permute broke contiguity.
3. Block 3: gather allocates. The cast doubles hidden. It is not free.

Stretch

Rewrite Block 1 so the only full copies are the ones you can justify in a review. List each remaining allocation in one clause.

Check your work

Block 1. x: (16, 1024, 4096) fp16 = 16 × 1024 × 4096 × 2 = 134217728 bytes = 128 MiB.

a = transpose(1,2): shape (16, 4096, 1024), same storage, not contiguous.

b = contiguous(): 128 MiB copy. First unavoidable full copy if anything later needs dense row-major.

c = reshape: view of b (now contiguous).

d = to(fp32): 256 MiB new storage. Traffic ≈ 128 MiB read + 256 MiB write.

e = d[:, :512]: view. Shape (65536, 512), still holding all of d alive.

f = clone(): 65536 × 512 × 4 = 128 MiB.

Live if nothing freed: x 128 + b 128 + d 256 + f 128 = 640 MiB. (If you dropped x after b, 512 MiB.) Downstream could take fp16 and skip d/f.

Block 2. w.t() view. GEMM reads it with whatever layout the backend accepts (may transpose internally — that is the vendor's copy, not yours). y is (8, 4096) bf16 = 64 KiB, allocates. z view. p view (strides swapped). q = reshape on a non-contiguous permute typically copies 64 KiB.

Block 3. emb = 32000 × 4096 × 2 ≈ 250 MiB. emb[ids] is an index/gather, not a view: it allocates (4, 2048, 4096) bf16 = 64 MiB. .float() allocates 128 MiB and doubles traffic for that tensor. Same shape, twice the bus.

Debrief

Shape is a story. Storage is the bill. If you cannot fill the table, you are not ready to profile the block.