← All questions
Inference & ServingVery Hard· Attention Mechanism· Compute Intensity· Hardware Architecture

At what batch size does an attention kernel transition...

The crossover from memory-bound to compute-bound occurs when the FLOPs per attention step (approximately 2 * B * T * d) divided by the bytes loaded (approximately 2 * B * T * d + B * d^2) exceeds the hardware’s ridge point, which is roughly 300 for an H100 GPU. For small head dimensions, the crossover batch size is very large and difficult to reach. Conversely, for larger head dimensions, the crossover is reached much earlier. Techniques like Flash Attention are designed to mitigate this by fusing tiles to reduce HBM round trips, regardless of the specific batch size.

Code playground (Python)