GPU benchmarks, 70B memory sizing and AI compilers

Modern AI systems engineering brings together kernel-level benchmarking, capacity and bandwidth math, and multi-stage compiler lowering. Below is a complete breakdown of all three areas.

1. Complete Benchmark Script: Eager PyTorch vs. torch.compile vs. Custom Triton

Because GPU execution is asynchronous, accurate CUDA benchmarking requires warmup passes (to trigger JIT compilation and allocate memory caches) and torch.cuda.synchronize() (or CUDA Events) before reading wall-clock timers.

Below is a standalone Python script comparing Eager PyTorch, torch.compile (TorchInductor), and a Custom Triton Kernel on a fused row-wise Softmax workload across a 4096 × 8192 FP16 matrix.

import torch
import triton
import triton.language as tl

# 1. Custom Fused Triton Kernel
@triton.jit
def fused_softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr
):
    row_idx = tl.program_id(0)
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # Single read from HBM into SRAM/registers
    row = tl.load(row_start_ptr + col_offsets, mask=mask, other=-float('inf')).to(tl.float32)
    row_minus_max = row - tl.max(row, axis=0)
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = (numerator / denominator).to(tl.float16)

    # Single write back to HBM
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    tl.store(output_row_start_ptr + col_offsets, softmax_output, mask=mask)

def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x)
    fused_softmax_kernel[(n_rows,)](
        y, x, x.stride(0), y.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE
    )
    return y

# 2. Unfused Eager PyTorch Baseline (Forces multiple HBM round-trips)
def naive_pytorch_softmax(x: torch.Tensor) -> torch.Tensor:
    x_fp32 = x.float()
    x_max = x_fp32.max(dim=-1, keepdim=True)[0]
    exp_x = torch.exp(x_fp32 - x_max)
    return (exp_x / exp_x.sum(dim=-1, keepdim=True)).half()

# 3. TorchInductor Compiled Function
compiled_softmax = torch.compile(naive_pytorch_softmax, mode="max-autotune")

# 4. Accurate CUDA Event Benchmarking Harness
def benchmark_fn(fn, x, warmup=25, rep=100):
    for _ in range(warmup):
        fn(x)
    torch.cuda.synchronize()

    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(rep)]
    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(rep)]

    for i in range(rep):
        start_events[i].record()
        fn(x)
        end_events[i].record()

    torch.cuda.synchronize()
    times_ms = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]
    median_ms = sorted(times_ms)[rep // 2]

    # Effective Memory Bandwidth (GB/s): 1 read + 1 write of FP16 tensor
    bytes_transferred = 2 * x.numel() * x.element_size()
    gb_per_sec = (bytes_transferred / (median_ms * 1e-3)) / 1e9
    return median_ms, gb_per_sec

if __name__ == "__main__":
    assert torch.cuda.is_available(), "CUDA GPU required."
    x = torch.randn((4096, 8192), device="cuda", dtype=torch.float16)

    # Verify numerical accuracy against native PyTorch
    ref = torch.softmax(x, dim=-1)
    assert torch.allclose(triton_softmax(x), ref, atol=1e-2)
    assert torch.allclose(compiled_softmax(x), ref, atol=1e-2)

    for name, fn in [
        ("Naive Eager PyTorch", naive_pytorch_softmax),
        ("torch.compile (Inductor)", compiled_softmax),
        ("Custom Triton Kernel", triton_softmax),
    ]:
        ms, bw = benchmark_fn(fn, x)
        print(f"{name:<26} | Latency: {ms:.4f} ms | Effective BW: {bw:.1f} GB/s")

• What the Results Show: On an H100 (3,350 GB/s peak HBM bandwidth), Naive Eager PyTorch typically achieves only ~600–800 GB/s effective bandwidth because it launches 5 separate kernels and writes intermediate tensors to VRAM repeatedly. Both torch.compile and the Custom Triton Kernel fuse the entire chain into a single pass, routinely hitting 2,600–3,000+ GB/s (80–90% of theoretical hardware peak).

2. Hardware Sizing Plan: Serving a 70B LLM in Production

To size a cluster accurately, you must calculate static weight memory, dynamic KV cache memory, and memory-bandwidth-bound decode throughput.

Target Workload Assumptions

• Model: Llama-3-style 70B Dense Model (80 layers, 8 KV heads via GQA, 128 head dimension).

• Precision: FP8 weights (1 byte per param) and FP8 KV Cache (1 byte per element).

• Concurrency & Context: 64 concurrent users, each with an 8,192-token context window.

• Hardware Candidate: NVIDIA H100 80GB SXM (80 GB HBM3, 3.35 TB/s memory bandwidth).

Step 1: Calculate Static Model Weight VRAM

Weight VRAM=70× 10^(9) params× 1 byte (FP8)≈ 70 GB

Step 2: Calculate Dynamic KV Cache VRAM

For every token, each layer stores a Key vector and a Value vector across 8 GQA heads of dimension 128:

KV Bytes per Token = 2 (K & V) × 80 layers × 8 KV heads × 128 dim × 1 byte = 163,840 bytes ≈ 160 KB/token

Across 64 concurrent users at 8,192 tokens each (524,288 total active tokens):

Total KV Cache VRAM=524,288 tokens× 163,840 bytes≈ 85.9 GB

Step 3: Determine GPU Count & Tensor Parallelism

Total Required VRAM=70 GB (Weights)+85.9 GB (KV Cache)+~15 GB (CUDA Graphs/Activations)≈ 171 GB

• Minimum Hardware: 4× H100 80GB GPUs connected via NVLink (320 GB total HBM), running with Tensor Parallelism = 4 (TP=4).

• Per-GPU Memory Footprint: Each GPU holds 17.5 GB of sharded weights + 21.5 GB of KV cache + ~4 GB runtime overhead = ~43 GB used out of 80 GB, leaving generous headroom for traffic spikes.

Step 4: Estimate Token Generation Throughput (Decode Speed)

During decode, generating one step across all 64 users requires reading all 70 GB of model weights plus the 85.9 GB KV cache once from HBM (155.9 GB total read per step) across the 4 GPUs' combined memory bandwidth (4 × 3.35 TB/s = 13.4 TB/s, assuming ~75% bus efficiency = 10.05 TB/s effective):

Time per Decode Step (TPOT)≈ (155.9 GB)/(10,050 GB/s)≈ 15.5 ms per token

Per-User Speed≈ (1000 ms)/(15.5 ms)≈ 64.5 tokens/sec/user | Total System Throughput=64× 64.5≈ 4,128 tokens/sec

3. AI Compiler Stacks: Lowering Python Down to PTX Assembly

Modern AI compilers transform high-level dynamic Python code into bare-metal GPU machine instructions through a multi-stage Intermediate Representation (IR) pipeline.

• Why MLIR (Multi-Level Intermediate Representation) Dominates Modern AI Compilers.

• Historically, compilers like standard LLVM operated only on low-level scalar loops and pointers, losing sight of high-level matrix structure.
• MLIR allows compilers (such as Triton, XLA, and Mojo) to maintain multiple levels of abstraction simultaneously—reasoning about 2D matrix tiles and shared memory layouts at higher dialects before progressively lowering down to LLVM IR, PTX (Parallel Thread Execution) virtual assembly, and finally SASS (the exact binary machine code executed by your specific NVIDIA GPU generation).