GPU Architecture › Code Implementation

import torch

# Check if CUDA is available and inspect device properties
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

if torch.cuda.is_available():
    props = torch.cuda.get_device_properties(0)
    total_memory_gb = props.total_memory / 1024**3
    print(f"GPU: {props.name}")
    print(f"Total memory: {total_memory_gb:.1f} GB")
    print(f"SM count: {props.multi_processor_count}")
    print(f"Max threads per SM: {props.max_threads_per_multi_processor}")
    print(f"Max threads per block: {props.max_threads_per_block}")
    print(f"Warp size: {props.warp_size}")
    print(f"Major/minor compute capability: {props.major}.{props.minor}")

import torch


def memory_of_tensor(shape, dtype=torch.float32):
    """Return memory in megabytes for a tensor of given shape and dtype."""
    dtype_bytes = {
        torch.float32: 4,
        torch.float16: 2,
        torch.bfloat16: 2,
        torch.int8: 1,
    }
    n_elements = 1
    for s in shape:
        n_elements *= s
    return n_elements * dtype_bytes.get(dtype, 4) / 1024**2


# Components of a LLaMA-style transformer layer
d_model = 4096  # hidden dimension
n_heads = 32  # attention heads
d_ff = 16384  # FFN hidden dimension (4x d_model)

# Attention weight matrices (Q, K, V, O projections)
qkv_shape = (d_model, d_model)
attn_params = 4 * memory_of_tensor(qkv_shape, torch.bfloat16)

# FFN weight matrices (two linear layers with optional gating)
ffn_w1_shape = (d_model, d_ff)
ffn_w2_shape = (d_ff, d_model)
ffn_params = memory_of_tensor(ffn_w1_shape, torch.bfloat16) + memory_of_tensor(
    ffn_w2_shape, torch.bfloat16
)

total_layer = attn_params + ffn_params

import time

import torch


def benchmark_matmul(m, k, n, dtype=torch.float16, warmup=5, iters=20):
    """Benchmark matrix multiplication and return achieved TFLOP/s."""
    if not torch.cuda.is_available():
        return None
    a = torch.randn(m, k, dtype=dtype, device="cuda")
    b = torch.randn(k, n, dtype=dtype, device="cuda")

    # Warm up to ensure JIT compilation and device initialization are complete
    for _ in range(warmup):
        c = torch.mm(a, b)
    torch.cuda.synchronize()

    start = time.perf_counter()
    for _ in range(iters):
        c = torch.mm(a, b)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

    avg_time = elapsed / iters
    flops = 2 * m * k * n  # each multiply-add counts as 2 FLOPs
    tflops = flops / avg_time / 1e12
    return tflops


# Test across matrix sizes from small (inefficient) to large (near-peak)
sizes = [256, 512, 1024, 2048, 4096, 8192]
results_fp16 = []
results_fp32 = []

for s in sizes:
    tf16 = benchmark_matmul(s, s, s, dtype=torch.float16)
    tf32 = benchmark_matmul(s, s, s, dtype=torch.float32)
    results_fp16.append(tf16)
    results_fp32.append(tf32)

import time

import torch


def benchmark_op(fn, warmup=5, iters=20):
    """Time a GPU operation and return milliseconds per iteration."""
    if not torch.cuda.is_available():
        return None
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) / iters * 1000


n = 4096
device_str = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(n, n, dtype=torch.float16, device=device_str)
a = torch.randn(n, n, dtype=torch.float16, device=device_str)
b = torch.randn(n, n, dtype=torch.float16, device=device_str)

# Memory-bound: elementwise ReLU reads and writes the full matrix
relu_time = benchmark_op(lambda: torch.relu(x))

# Compute-bound: matrix multiplication has very high arithmetic intensity
matmul_time = benchmark_op(lambda: torch.mm(a, b))

# Memory-bound: softmax involves multiple passes over the data
softmax_time = benchmark_op(lambda: torch.softmax(x, dim=-1))

import torch
import torch.nn as nn


def get_gpu_memory_mb():
    """Return allocated and reserved GPU memory in MB."""
    if not torch.cuda.is_available():
        return 0, 0
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    return allocated, reserved


# Build a small transformer-like model for demonstration
class SimpleTransformerLayer(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ff1 = nn.Linear(d_model, d_ff)
        self.ff2 = nn.Linear(d_ff, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.gelu = nn.GELU()

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        ff_out = self.ff2(self.gelu(self.ff1(x)))
        return self.norm2(x + ff_out)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleTransformerLayer(d_model=512, n_heads=8, d_ff=2048).to(device)

# Measure memory at each stage of a forward pass
torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None

alloc_0, res_0 = get_gpu_memory_mb()

# Model weights
param_mb = (
    sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2
)

# Create input batch
batch_size, seq_len = 16, 256
x = torch.randn(batch_size, seq_len, 512, device=device, dtype=torch.float32)
alloc_1, res_1 = get_gpu_memory_mb()

# Forward pass (activations are created and stored)
y = model(x)
alloc_2, res_2 = get_gpu_memory_mb()

# Backward pass (gradients computed)
loss = y.mean()
loss.backward()
alloc_3, res_3 = get_gpu_memory_mb()