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

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.¹

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.

A tensor is a typed array with a shape and a device

Here are the names used for different array dimensions:

• A scalar is a single number. Rank 0.
• A vector is a list of numbers. Rank 1.
• A matrix is a grid of numbers. Rank 2.
• A tensor is the general word for any of these, including 3-D and higher.

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):

#!/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.

Reading rank, shape, dtype, and device

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.

Byte size is Week 4's formula, now measured

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 are the same size but not the same numbers

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:

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.

Element-wise operations and broadcasting

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.

Device: CPU tensor versus CUDA tensor

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.

Precision changes speed, not just memory

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:

#!/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.

Results

Here is Week 5 in one place:

What surprised me

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.

Mistakes and troubleshooting

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.

Production implications

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.

What I will learn next

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.²

Run it yourself

The public Week 5 lab has the tensor-basics script, the matmul benchmark, the
captured results, observations, and troubleshooting notes.³

1. CUDA track (parallel), Level 1:
   https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/cuda-track.md ↩

2. Week 6 roadmap:
   https://github.com/dramasamy/from-api-to-gpu/blob/main/roadmap/week-06.md ↩

3. Week 5 companion lab:
   https://github.com/dramasamy/from-api-to-gpu/tree/main/week-05-pytorch-tensors ↩