when your inference kernels are (theoretically) memory bound vs compute bound

KV caching

If Wk, Wv ∈ ℝdmodel × dmodel and bt ∈ ℝ1 × dmodel, then FLOPS become

2 · 2 · nlayers · dmodel2
  • the first 2: one MAC is two individual operations (multiply + add)
  • the second 2: Wk and Wv
  • dmodel2: the dot product

That was for one token. Now imagine storing KV cache for the entirety of batch + sequence:

(2 · 2 · nlayers · dmodel2) · nbatch · nsequence

There are two segments to generation in LLMs: prefill and decode. Prefill → compute bound. Decode → memory bound.

min-step during decode

Let's say during decode, we are trying to find min-step — the lower bound for time taken per generation of a token.

min-step = (batch size · KV cache size + parameters) / total memory bandwidth

Then max tokens/sec becomes

B / min-step = B / ((B · KV + P) / W) = (B · W) / (B · KV + P)
however, we know that KV cache during decode for one token is memory bound. but the MLP?

a more general min-step

min-step = (B · KV) / W + max(P / W, 2 · P · B / C)
  • (B · KV) / W: attention. always memory bound.
  • P / W: MLP memory
  • 2 · P · B / C: MLP compute (FLOPS/sec)

If W = C, then as P increases, you become compute bound. But more importantly, the MLP can be either compute or memory bound based on B — if B increases, you become increasingly closer to roofline and eventually compute bound, since you are increasing computation relative to B compared to always bringing just P params regardless of batch size B.

critical batch size

Bcrit = β · α
  • β = (bytes per param) / (bytes per activation) → memory / compute
  • α is hardware-defined

If you increase bytes per param, you increase the bytes you need to transfer by that rate. If you increase bytes per activation, you are increasing the width you are working with, so n times the number of operations.

B ↓ → higher throughput (generally), lower GPU utilization.

example

  • B = 4
  • P = 30B
  • 16 chips
  • bandwidth = 8.2 · 1011 bytes/sec
  • 1.97 · 1014 FLOPS/sec
  • 8192 context
  • 100 kB / token KV cache
min-step = (4 · 8192 · 100 · 103) / (8.2 · 1011 · 16) + (30 · 109) / (8.2 · 1011 · 16)
         = 2.54 ms
here you can assume memory bound because B is low. throughput is fast, but again, memory bound.

B = 256

Try with B = 256, larger B. Then

min-step = (256 · 100 · 103 · 8192) / (8.2 · 1011 · 16) + (2 · 30 · 109 · 256) / (1.97 · 1014 · 16)
         = 20.96 ms

As you can see, throughput and time taken is better in this case for the respective measurements, and GPU utilization is probably better.

loading a 52B model

Say we are loading a 52B parameter model on device. First, with dtype = float16,

2 · 52 · 109 = 1.04 · 1011 bytes = 104 GB

So we would need at least 3 GPUs to load this. How much do we have left for the KV cache?

120 − 104 = 16 GB

KV cache size (per token):

2 · 2 · nlayers · nheads · dhead = 4 · 64 · 8192 = 2,097,152 bytes = 0.002 GB

or equivalently

KV cache size = 2 · (bytes per element) · nlayers · nheads · dhead
              = 2 · (bytes per element) · nlayers · dmodel

This means 16 GB / 0.002 GB = 8000 tokens of capacity left.

In that case, we could probably do 4 batches with 2048 tokens each, but that still becomes memory bound since Bcrit = 208 in this case. We ideally want each GPU involved to be hitting close to Bcrit (all in total).

GQA

To reduce the impact of storage of the KV cache, we want to reduce the amount of memory allocated for it. Something like GQA works: we share KV caches across heads, so if we have a group of 8, then we reduce memory size theoretically by a factor of 8.

in this case, 2 Q heads attend over the same KV heads, so we need half as much memory, but every pair of Q heads would be sharing KV heads and information.
← back