Zachary Montague

Lab

One request through LLM inference

One prompt, followed from text to tensors to streamed tokens, on two GPUs with tensor parallelism.

Illustrative. Generic shapes, not measurements.

seq is the number of tokens, hidden is the model width, and vocab is the vocabulary size.

  1. 1 Tokenize. The prompt is split into tokens, and each token maps to an integer ID in the vocabulary.
  2. 2 Embed. Each ID selects one row of the embedding table. The rows stack into a tensor of shape [seq, hidden].
  3. 3 Shard across two GPUs. The weights are split across 2x H100 with TP=2. Each GPU holds half of every layer.
  4. 4 Run the layers. Every layer runs attention, then an MLP. Each GPU computes a partial result, and an all-reduce adds the two halves so both GPUs hold the full activations.
  5. 5 Prefill the KV cache. During prefill the whole prompt runs in one pass. Keys and values for every position are written to the KV cache on each GPU.
  6. 6 Sample the next token. The final layer produces logits of shape [vocab]. Softmax turns them into probabilities, and sampling picks the next token.
  7. 7 Decode and stream. Decode runs one token at a time. Each step reads the KV cache, appends one new entry, and streams the token back to the user.