From API to GPU, Week 5: Tensors, the Data Structure Behind Every Model

A developer's ongoing 32-week project to understand AI inference reached Phase 2, focusing on tensors, the core data structure of machine learning models. The developer demonstrated creating tensors in PyTorch, inspecting their properties, and moving them to a GPU, while also starting a parallel CUDA track. The work runs on an NVIDIA DGX Spark system.

Phase 2 of 8: Enough ML to understand inference. Week 5 of 32. Phase 1 was about running models. Phase 2 is about understanding what happens inside them, starting with the one data structure they are all built from: the tensor . This week I stop talking about model files and start touching the actual numbers, in PyTorch, on the GPU. If you write software, a tensor is a typed, multi-dimensional array that lives on a specific device. By the end of this post I can create tensors, read their shape and byte size, move them to the GPU, and measure how precision changes both speed and memory. This is also where the parallel CUDA track starts. I am not writing GPU kernels yet, the small programs that run on the GPU. Level 1 is just being a competent CUDA user: picking a device, moving data to it, and timing GPU work correctly.1 Everything runs on the DGX Spark over ssh spark . I reuse the Week 1 PyTorch environment, ~/venvs/w1 , which already has a CUDA build of PyTorch. Here are the names used for different array dimensions: Rank is the number of dimensions. Shape is the size along each dimension. Dtype is the number format, the same FP32, FP16, and BF16 I measured in Week 4. Device is where the tensor lives, the CPU or the GPU. Instead of describing this, I print it. The first script builds a scalar, vector, matrix, and 3-D tensor and reports each one's rank, shape, dtype, byte size, and device. It then shows three operations I explain right after: an element-wise add position by position , a broadcast a smaller tensor applied across a larger one , and a host-to-device transfer moving a tensor to the GPU : bash /usr/bin/env python3 """Week 5 - tensor basics: rank, shape, dtype, bytes, and device. Shows a scalar, vector, matrix, and 3D tensor, then one element-wise operation, one broadcast, and one CPU-to-GPU transfer. Every printed size and byte count comes from PyTorch, not from a hand estimate. """ from future import annotations import torch def describe name: str, tensor: torch.Tensor - None: print f"{name:8} rank={tensor.ndim} " f"shape={tuple tensor.shape } " f"dtype={str tensor.dtype .replace 'torch.', '' } " f"elem={tensor.element size }B " f"total={tensor.nbytes}B " f"dev={tensor.device}" def main - None: scalar = torch.tensor 3.0 vector = torch.tensor 1.0, 2.0, 3.0 matrix = torch.zeros 2, 3 tensor3d = torch.zeros 2, 3, 4 print "=== rank, shape, dtype, bytes, device ===" describe "scalar", scalar describe "vector", vector describe "matrix", matrix describe "tensor3d", tensor3d print "\n=== same shape, three precisions ===" for dtype in torch.float32, torch.float16, torch.bfloat16 : describe str dtype .replace "torch.", "" , torch.zeros 1024, 1024, dtype=dtype print "\n=== element-wise operation ===" a = torch.tensor 1.0, 2.0, 3.0 b = torch.tensor 10.0, 20.0, 30.0 print "a + b =", a + b .tolist print "\n=== broadcasting ===" matrix = torch.ones 2, 3 row = torch.tensor 1.0, 2.0, 3.0 print "matrix shape", tuple matrix.shape , "+ row shape", tuple row.shape print "result:\n", matrix + row .tolist print "\n=== device transfer ===" if torch.cuda.is available : cpu tensor = torch.ones 3 gpu tensor = cpu tensor.to "cuda" print "cpu tensor.device", cpu tensor.device print "gpu tensor.device", gpu tensor.device else: print "CUDA not available; skipping GPU transfer" if name == " main ": main I ran it on the Spark. These commands use public/... paths because I run them from the parent repository. If you cloned the public companion repo, drop the public/ prefix and run from that repo root. ssh spark '~/venvs/w1/bin/python -' \ < public/week-05-pytorch-tensors/tensor basics.py === rank, shape, dtype, bytes, device === scalar rank=0 shape= dtype=float32 elem=4B total=4B dev=cpu vector rank=1 shape= 3, dtype=float32 elem=4B total=12B dev=cpu matrix rank=2 shape= 2, 3 dtype=float32 elem=4B total=24B dev=cpu tensor3d rank=3 shape= 2, 3, 4 dtype=float32 elem=4B total=96B dev=cpu === same shape, three precisions === float32 rank=2 shape= 1024, 1024 dtype=float32 elem=4B total=4194304B dev=cpu float16 rank=2 shape= 1024, 1024 dtype=float16 elem=2B total=2097152B dev=cpu bfloat16 rank=2 shape= 1024, 1024 dtype=bfloat16 elem=2B total=2097152B dev=cpu === element-wise operation === a + b = 11.0, 22.0, 33.0 === broadcasting === matrix shape 2, 3 + row shape 3, result: 2.0, 3.0, 4.0 , 2.0, 3.0, 4.0 === device transfer === cpu tensor.device cpu gpu tensor.device cuda:0 There is a lot in that output, so here is what matters. The shape notation is worth decoding once. means zero dimensions, a single number. 3, means one dimension holding three values, the trailing comma marking it as a shape rather than a plain number. 2, 3 means two rows and three columns. 2, 3, 4 adds a third dimension. I never set a dtype in those constructors, and every tensor still came out float32 . That is PyTorch's default floating-point type. The precision block below sets the dtype on purpose. Look at the vector: shape 3, , 4 bytes per element, 12 total bytes. That is just 3 times 4. The matrix is 2 times 3 times 4, which is 24 bytes. The rule is the same one from Week 4: total bytes equals element count times bytes per element. PyTorch reports it directly with nbytes , so I no longer have to estimate. The three-precision block makes the point sharper. The same 1024, 1024 shape is 4,194,304 bytes at FP32 but 2,097,152 bytes at FP16 and BF16, exactly half. A dtype decides two things at once: the byte size of every number, and which values the number can represent. This experiment measures the byte size. The next one looks at the values. FP16 and BF16 both report 2 bytes, so they use the same memory. The difference is how they split those 16 bits between range how large a number can get and step size how finely close numbers can be told apart . I read both directly with torch.finfo : python ssh spark '~/venvs/w1/bin/python - <<PY import torch for name, dt in "fp32", torch.float32 , "fp16", torch.float16 , "bf16", torch.bfloat16 : fi = torch.finfo dt print f"{name:5} max={fi.max:.3e} smallest step near 1={fi.eps:.3e}" PY' fp32 max=3.403e+38 smallest step near 1=1.192e-07 fp16 max=6.550e+04 smallest step near 1=9.766e-04 bf16 max=3.390e+38 smallest step near 1=7.812e-03 FP16 tops out at about 65,500 but has the finer step. BF16 keeps almost FP32's huge range but has a coarser step. That is the tradeoff in one line: FP16 gives finer detail over a small range, BF16 gives a wide range with coarser detail. Which one is better for training and inference is a later-week topic. This week I only need to see that same size does not mean same numbers. Two everyday operations show up constantly, so they are worth naming. An element-wise operation works position by position. Adding 1, 2, 3 and 10, 20, 30 gives 11, 22, 33 . Nothing mixes across positions. Broadcasting lets PyTorch treat compatible dimensions as repeated, without making a full copy. Here the row's length 3 matches the matrix's last dimension of 3, so PyTorch adds the row to every row of the matrix. The output shows 1, 2, 3 added to both rows of ones, giving 2, 3, 4 twice. This is how model code adds the same bias vector to every row of a result in one line. The last block shows the CUDA-track idea in miniature. A tensor's device is where PyTorch will run its operations, either the host the CPU side or a CUDA GPU. cpu tensor.device is cpu . After .to "cuda" , gpu tensor.device is cuda:0 . That .to "cuda" is a host-to-device transfer , and cuda:0 is device selection , naming the first GPU. On the DGX Spark the CPU and GPU share one physical memory pool, so this is not a copy across a separate VRAM. The output proves the device label changed to cuda:0 , which is what the GPU needs before it will compute on the tensor. It does not measure transfer cost. This week begins CUDA-track Level 1: pick a device, put your data on it, measure correctly, and know what happens on out-of-memory. Level 1 continues over the next few weeks. Week 1 already covered the driver, toolkit, and CUDA-enabled PyTorch underneath it. Week 4 showed precision changes storage. Week 1 showed the GPU running a large matrix multiply much faster than the CPU. Week 5 puts those together and adds the missing dimension: precision also changes how fast the same operation runs. The benchmark multiplies two square matrices many times per precision. Two details keep GPU timing honest, both from the CUDA track. First a warmup : the first GPU call pays a one-time setup cost, so I run several throwaway iterations before timing. Second synchronization : CUDA runs work asynchronously, so without torch.cuda.synchronize I would be timing how long it takes to queue the work, not to finish it. Here is the script: bash /usr/bin/env python3 """Week 5 - measure matmul speed and memory for FP32, FP16, and BF16. Times a square matrix multiply on the GPU for three precisions with warmup and CUDA synchronization, then compares one CPU run with one GPU run at FP32. All timings use time.perf counter around synchronized GPU work. Results print as a table and optionally save to JSON. """ from future import annotations import argparse import json import statistics import time from pathlib import Path import torch DTYPES = { "fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, } def parse args - argparse.Namespace: parser = argparse.ArgumentParser description= doc parser.add argument "--size", type=int, default=4096 parser.add argument "--iters", type=int, default=20 parser.add argument "--warmup", type=int, default=5 parser.add argument "--output", type=Path return parser.parse args def time matmul a: torch.Tensor, b: torch.Tensor, iters: int, warmup: int - float: cuda = a.device.type == "cuda" for in range warmup : a @ b if cuda: torch.cuda.synchronize samples = for in range iters : start = time.perf counter a @ b if cuda: torch.cuda.synchronize samples.append time.perf counter - start return statistics.median samples def gflops size: int, seconds: float - float: return 2 size 3 / seconds / 1e9 def main - None: args = parse args size = args.size has cuda = torch.cuda.is available device = "cuda" if has cuda else "cpu" rows = print f"matmul {size}x{size}, iters={args.iters}, device={device}\n" print f"{'precision':10} {'median ms': 10} {'gflops': 10} " f"{'peak mib': 10}" for name, dtype in DTYPES.items : if not has cuda: continue torch.cuda.empty cache torch.cuda.reset peak memory stats a = torch.randn size, size, device=device, dtype=dtype b = torch.randn size, size, device=device, dtype=dtype seconds = time matmul a, b, args.iters, args.warmup peak mib = torch.cuda.max memory allocated / 1024 2 row = { "precision": name, "median ms": round seconds 1000, 3 , "gflops": round gflops size, seconds , 1 , "peak mib": round peak mib, 1 , } rows.append row print f"{name:10} {row 'median ms' : 10.3f} {row 'gflops' : 10.1f} " f"{row 'peak mib' : 10.1f}" del a, b print "\n=== CPU versus GPU at FP32 ===" a cpu = torch.randn size, size b cpu = torch.randn size, size cpu iters = max 3, args.iters // 4 cpu s = time matmul a cpu, b cpu, cpu iters, 1 print f"cpu fp32 {cpu s 1000:10.3f} ms {gflops size, cpu s :8.1f} " f"gflops" speedup = None if has cuda: gpu fp32 = next r for r in rows if r "precision" == "fp32" speedup = round cpu s 1000 / gpu fp32 "median ms" , 1 print f"gpu fp32 {gpu fp32 'median ms' :10.3f} ms " f"{gpu fp32 'gflops' :8.1f} gflops" print f"gpu is {speedup}x faster than cpu at fp32" if args.output: payload = { "size": size, "warmup": args.warmup, "gpu iters": args.iters, "cpu iters": cpu iters, "device": device, "precisions": rows, "cpu fp32 ms": round cpu s 1000, 3 , "gpu speedup fp32": speedup, } args.output.write text json.dumps payload, indent=2 + "\n" if name == " main ": main gflops reports billions of floating-point operations per second. A square matmul of size n produces n^2 output cells, and each cell does about n multiplications and n additions, so the total is about 2 n^3 operations. Dividing that by the seconds taken gives a speed I can compare across precisions. The GPU results below are the median of 20 timed iterations after warmup; the CPU result is the median of 5, since each CPU run is much slower. I ran it at 4096: ssh spark '~/venvs/w1/bin/python - --size 4096 --iters 20' \ < public/week-05-pytorch-tensors/benchmark matmul.py matmul 4096x4096, iters=20, device=cuda precision median ms gflops peak mib fp32 7.440 18473.0 224.0 fp16 1.548 88807.9 128.0 bf16 1.576 87202.3 128.0 === CPU versus GPU at FP32 === cpu fp32 168.281 ms 816.7 gflops gpu fp32 7.440 ms 18473.0 gflops gpu is 22.6x faster than cpu at fp32 Two clear results. First, the CPU-versus-GPU gap. The CPU does the FP32 matmul in about 168 ms. The GB10 GPU does it in about 7.4 ms, roughly 22.6 times faster. That matches the 22.5x I measured in Week 1, a good sign the setup is consistent. Second, the new result: precision changes speed. FP16 and BF16 finished the same matmul in about 1.55 ms versus 7.44 ms for FP32, roughly 4.8 times faster, on the exact same shape and hardware, just by using 16-bit numbers. Memory dropped too. The two input matrices plus the temporary output matrix hold 192 MiB of payload at FP32 and 96 MiB at FP16 or BF16, exactly half. PyTorch's measured peak, which is the most memory it had allocated at once 1 MiB is 1024 x 1024 bytes , fell from 224 MiB to 128 MiB. That is about 43 percent, not exactly half, because the peak includes roughly 32 MiB of extra working allocation on top of the three matrices. FP16 and BF16 landed nearly equal in this run, about 1.55 versus 1.58 ms. They are the same size and run at nearly the same speed here, and they can swap order between runs. The choice between them is about the range-and-step tradeoff shown earlier, not about performance on this test. Here is Week 5 in one place: | Item | Verified value | |---|---| | FP32 bytes per element | 4 | | FP16 / BF16 bytes per element | 2 | | 1024x1024 FP32 tensor | 4,194,304 bytes | | 1024x1024 FP16 / BF16 tensor | 2,097,152 bytes | | GPU FP32 matmul 4096 | 7.44 ms, ~18,473 GFLOP/s | | GPU FP16 matmul 4096 | 1.55 ms, ~88,808 GFLOP/s | | GPU BF16 matmul 4096 | 1.58 ms, ~87,202 GFLOP/s | | CPU FP32 matmul 4096 | 168.28 ms, ~817 GFLOP/s | | GPU vs CPU at FP32 | about 22.6x faster | | FP16 vs FP32 speed | about 4.8x faster | | FP32 vs FP16 measured peak | 224 MiB vs 128 MiB | The byte sizes matching Week 4 exactly was reassuring rather than surprising: a 1024x1024 FP32 tensor really is 4,194,304 bytes, no rounding. The real surprise was that precision buys speed, not only memory. I already knew 16-bit numbers use half the bytes. Seeing the same matmul run almost five times faster in FP16 made it clear why production inference leans on lower precision. It saves memory and time at once, at least on this one operation. My first timing attempt forgot torch.cuda.synchronize . Because CUDA runs work asynchronously, I was timing how long it took to launch the matmul, not to finish it, so the GPU looked far faster than it really is. Adding a warmup and a synchronize per timed iteration fixed it. Week 1 shows the related effect: its cold first GPU run took much longer than its warmed-up runs. GPU work is asynchronous, so time it with a warmup and a sync or you time the wrong thing. Single timings are also noisy, so the benchmark reports the median of 20 GPU iterations and 5 CPU iterations. The absolute milliseconds shift between runs, and FP16 and BF16 can trade places, but the groupings hold: 16-bit beats FP32, and the GPU beats the CPU. This week explains a choice I will keep seeing: serving models in FP16 or BF16 instead of FP32. Lower precision halves the tensor payload and, on this hardware, ran this dense matmul almost five times faster. That is one operation, not a whole model. A real model also reads weights from memory and runs other steps, so it will not simply be 4.8 times faster end to end. Even so, the direction is why production inference leans on 16-bit precision: it saves both memory and time. The CUDA-track habits matter for the rest of the series too. Correct GPU timing, explicit device placement, and watching peak memory are the starting tools. They measure elapsed time and capacity, not utilization or bandwidth, so they cannot yet classify a workload as compute-bound or memory-bound. Later CUDA-track profiling adds those measurements. Week 6 steps up from single tensors to a tiny neural network: inputs, a linear layer, an activation, a loss, and a training loop small enough to watch the weights change. That connects this week's tensor operations to how a model actually learns.2 The public Week 5 lab has the tensor-basics script, the matmul benchmark, the captured results, observations, and troubleshooting notes.3 CUDA track parallel , Level 1: https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/cuda-track.md https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/cuda-track.md ↩ Week 6 roadmap: https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/week-06.md https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/week-06.md ↩ Week 5 companion lab: https://github.com/dramasamy/from-api-to-gpu/tree/main/week-05-pytorch-tensors https://github.com/dramasamy/from-api-to-gpu/tree/main/week-05-pytorch-tensors ↩