Inference: Batch Size, Memory and Precision

Lesson

Process exactly 256 examples through a two-layer projection. Try microbatches and optionally autocast to a supported dtype; always return float32 in original order. Compare peak allocated/reserved memory, samples/sec, total latency and numerical error against FP32. Reserved memory includes the allocator cache; incremental peak excludes already-live inputs, weights and reference output. A smaller batch can reduce temporary activations while adding launches. This is inference only: no backward, optimizer, gradient accumulation or checkpointing claims. The next lab separately validates training gradients.

Measurement contract: the real read-only driver defines WORKLOAD. Two warmups precede five paired trials with alternating baseline/candidate order. Each trial uses fresh deterministic inputs. Synchronized host time includes solve and completion, excludes input construction and reference verification. METRICS reports every sample, median and IQR; memory includes persistent buffers and allocator state. One additional solve is isolated in judged for Trace. No speedup is promised. Run the unchanged starter first.

What you will learn

• Identify the bottleneck from the appropriate evidence.
• Preserve the output contract on fresh inputs.
• Compare paired measurements and numerical error.

Starter code

    return baseline(case)