48. Profiling Two Implementations of the Same Op

Problem

The same mathematical operation can be written many ways in PyTorch, and they are not equally fast: a Python-level loop over batch/head dimensions computing one matmul at a time forces the Python interpreter and dozens of tiny kernel launches into the hot path, while a single vectorized call (batched matmul or einsum) lets one fused, well-optimized kernel handle the whole tensor. Before reaching for a custom kernel, an engineer's first move is almost always: implement the naive version, implement the vectorized version, confirm they agree numerically, then profile both to find (and justify) the bottleneck.

Given the naive loop-based implementation, implement the vectorized version and the profiling harness that compares them.

Source: src/48_profile_op_bottleneck.py

def profile_and_compare(
    q: torch.Tensor,
    k: torch.Tensor,
    num_runs: int = 10,
) -> dict:
    ...

Examples:
>>> q = torch.randn(2, 3, 8, 16)
>>> k = torch.randn(2, 3, 8, 16)
>>> result = profile_and_compare(q, k, num_runs=5)
>>> result["outputs_match"]
True

Step-by-Step Approach

1. Write the vectorized implementation with a single torch.einsum("bhld,bhmd->bhlm", q, k) call, computing every (batch, head) pair's score matrix in one dispatched kernel instead of a Python-level loop.
2. Warm up both the loop and vectorized paths once before timing — the first call pays one-time costs (allocator warmup, kernel autotuning) that would otherwise skew the measurement.
3. Time the loop implementation across num_runs repetitions using time.perf_counter() around the whole loop, then divide by num_runs for a per-call average.
4. Repeat the same timing procedure for the vectorized implementation.
5. Verify correctness with torch.allclose between the two implementations' outputs — a profiling result is meaningless if the two versions don't compute the same thing.
6. Return a dict with both average times in milliseconds and the correctness flag.

The key insight is that profiling only tells you something useful once correctness is nailed down first — comparing the wall-clock time of two implementations that don't actually agree numerically is comparing nothing.

Reference solution

def batched_scores_vectorized(q, k):
    # one einsum call handles every (batch, head) pair at once, dispatching
    # to a single batched-matmul kernel instead of B*H separate Python calls
    return torch.einsum("bhld,bhmd->bhlm", q, k)

def profile_and_compare(q, k, num_runs=10):
    # warm up: first calls pay allocator/kernel-autotune costs
    _ = batched_scores_loop(q, k)
    _ = batched_scores_vectorized(q, k)

    start = time.perf_counter()
    for _ in range(num_runs):
        out_loop = batched_scores_loop(q, k)
    loop_elapsed = time.perf_counter() - start

    start = time.perf_counter()
    for _ in range(num_runs):
        out_vec = batched_scores_vectorized(q, k)
    vec_elapsed = time.perf_counter() - start

    return {
        "loop_ms": loop_elapsed * 1000.0 / num_runs,
        "vectorized_ms": vec_elapsed * 1000.0 / num_runs,
        "outputs_match": torch.allclose(out_loop, out_vec, atol=1e-5),
    }

Key Functions & Tricks

• torch.einsum — expresses the batched matmul with an explicit index pattern, one dispatched kernel instead of a Python-level double loop.
• time.perf_counter() — the correct low-overhead, monotonic timer for microbenchmarks on CPU (versus time.time(), which isn't guaranteed monotonic).
• Warmup iterations — excluded from the timed region so one-time costs don't pollute the steady-state measurement.
• torch.allclose — the correctness gate that must pass before a timing comparison means anything.
• torch.profiler.profile / torch.cuda.Event — the heavier-weight real tools for GPU-accurate timing and per-op memory/time breakdowns, worth naming even though this problem uses simple CPU wall-clock timing for determinism.

How to Recognize This Pattern