This article was written with the assumption that the reader has some understanding of the transformer architecture of LLMs with basic arithmetic skills. I also encourage the reader to have some understanding of model parameter counting. Read [2] as Chen, C. advised in her transformer arithmetic article [1].

Model inferencing for thousands of users and let alone millions or billions requires a high level of craftsmanship to get right.

Doing this efficiently requires optimizing both the runtime and infrastructure layers of your inferencing stack. Runtime optimization cares about writing efficient kernels for fast attention computation, KV reuse, kernel fusing to minimise GPU bubbles etc. The infrastructure layer borders around efficiently scheduling inferencing workloads across a cluster of GPUs, dealing with distributed systems tradeoffs like fault tolerance, correctness analysis with formal verification etc.

In this article, we want to understand how communication and compute boundedness impact inferencing speed.

LLM inferencing comprises of two math intensive stages - prefill and decode that require as much compute, memory/communication bandwidth as possible to minimize latency and increase throughput. Majority of the time consumed during inferencing happens during the decode stage due to the nature of how matmuls are calculated. Prefill matmul computation is embarrassingly parallel because there isn’t a need for tokens to be related to each other in such a way that makes semantic sense to the user. In decode, matmuls are computed in a sequential manner i.e on a token by token basis because each output token has to be a consequence of previous tokens in other to maintain attention. The most important problem in inferencing then becomes “how do we speed up the decode stage?”

There are bunch of variables that affect inferencing speed and we would touch on them:

1. Parameter size: Model dimension directly correlates with the number of FLOPs it takes to produce an output token. It also influences the amount of time required to stream the weights to the GPU. A bigger model means longer inference time.
2. GPU VRAM/HBM: This is the factory for tokens. The bigger the flop count the faster the matmul calculation and the bigger the memory bandwidth the faster the GPU streams weights from the VRAM to SRAM.
3. Parallelism technique: Depending on the technique, the speed for inferencing can be greatly affect by the cadence at which tensors are between GPUs in a single node. NVIDIA uses NVLink as the RDMA fabric for inter-gpu comms (synchronizing activations which we will discuss in sufficient detail later on) so a B300 with about 1.8TB/s interconnect means that it can share tensors quickly between GPUs on physically separate chips (choosing a parallelism type comes with its own tradeoff).
4. Batch size: This basically specifies how much inference request can a GPU handle per forward pass.
5. GPU capacity for weights and kv cache: kv cache utilises the left over space after loading the HBM with model weights. Having enough space to store kv cache is the difference between O(N) and O(N^2) time complexity during decode.

A fine-balance configuration is needed between these variables in other to achieve maximum token per watt for your GPUs.

KV cache

In a standard transformer, size of a token in bytes after one forward pass is equivalent to:

$$ 2~.~2~.~n_{layers}~.~n_{heads}~.~d_{head} \tag 1 $$

One of the “2” represents two bytes that’s required for an FP16 value while the second “2” represents the fact that k and v tokens are stored.

k and v projections:

$$ W_k, W_v \in \mathbb{R}^{d_{model} \times d_{model}} \tag 2 $$

token embedding is a 1-d vector from the embedding matrix:

$$ t \in \mathbb{R}^{1 \times d_{model}} \tag 3 $$