Stage the Batch: One Copy, One Matmul

This function scores a batch of CPU feature rows against a weight matrix on the GPU - logits = data @ weight. The starter copies one row to the device inside a loop and runs a tiny matmul per row, so both the PCIe bus and the launch queue set the pace instead of the math. Keep the same result and stage the whole batch to the GPU once, then do a single batched matmul. GPUFlight counts host-to-device transfers AND kernel launches in the judged region, so it rewards staging the input once and batching the matmul.

Languages: PyTorch

Lesson

Moving a CPU tensor to the GPU one row at a time issues one host-to-device (H2D) transfer per iteration, and a per-row matmul issues one small GEMV launch per iteration - both cross expensive boundaries (the PCIe bus, the host launch queue) that dwarf the trivial per-row math. Staging the whole batch with a single `data.to('cuda')` moves the same bytes in one transfer, and a single batched matmul `data @ weight` runs one large GEMM instead of N tiny GEMVs. Same result, one transfer and a small number of launches instead of N of each. cuBLAS is far more efficient on one [N, F] x [F, C] GEMM than on N separate [F] x [F, C] GEMVs.

One batched matmul is a framework operation, not a guaranteed single GPU kernel. Inspect actual backend dispatch in the timeline. Return CUDA float32 [128,256].

What you will learn

Recognize a per-row `.to('cuda')` as a per-iteration PCIe transfer and a per-row matmul as a per-iteration kernel launch.
Stage an entire batch to the device once and keep the computation device-resident.
Replace a Python loop of tiny GEMVs with a single batched matmul (one GEMM).
Read host-to-device transfer count and kernel launch count together as signals of un-batched host code.

Starter code

    # Naive: move each row to the GPU inside the loop, then a tiny matmul per
    # row. One H2D transfer AND one small GEMV every iteration - the PCIe bus and
    # the launch queue set the pace, not the math.
    n = data.shape[0]
    out = torch.empty(n, weight.shape[1], device='cuda')
    for i in range(n):
        row = data[i].to('cuda')
        out[i] = row @ weight
    return out