Prefill and Decode Analysis with the Roofline Model
The problem, from first principles
Whether quantization (or any optimization) actually speeds up a workload depends on whether that workload is compute-bound or memory-bound — optimizing the wrong resource does nothing. Given GPU compute (FLOP/s) and memory bandwidth (bytes/s) specs, and a model/workload's FLOPs and bytes moved, classify a prefill step and a decode step as compute-bound or memory-bound using the roofline model, and estimate each step's wall-clock time.
From theory to code
intensity = flops / bytes_moved
ridge_point = peak_flops_per_sec / peak_bytes_per_sec
time = max(flops / peak_flops_per_sec, bytes_moved / peak_bytes_per_sec)
Constraints
- Compute arithmetic intensity, ridge point, both individual time estimates, the overall predicted time (max of the two), and a
'compute_bound'vs'memory_bound'classification. - Do this independently for a
'prefill'step and a'decode'step given separately, in one call. - Handle the edge case
bytes_moved == 0(pure-compute op, always compute-bound) without dividing by zero.
Hints
Hint: You don't need intensity to get the time right
Compute both flops / peak_flops_per_sec and bytes_moved / peak_bytes_per_sec directly — you don't strictly need the intensity/ridge-point comparison to get the right time, but it's the cleanest way to also report WHICH bound is active. The classification is just whichever of the two times is larger.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.