Writing an LLM inference loop

When you write a simple prompt in claude or chatgpt i always wondered what happens behind the scenes. For the sake of this article we are not going to be diving into the architecure of how all of this is setup. and the training of the model itslef.

Lets just dive into the loop that takes the input tokens, runs what we call inference i.e get the output from the input and what the model predicts and then outputs it on the screen.

Lets the goal of the model here. The model takes the tokens seen so far, runs attention over those context tokens build a representation of the current state, and then outputs logits: one raw score for each token in the vocabulary. Those logits represent how likely each vocabulary token is to be the next token.

Lets now divide this process into steps that require to produce an output token.

• Encode

  • Takes in the input text and converts it into token ids using a vocabulary tokneizer

  • For the words that do not have a corresponding token id we have some unknown token ids that can be returned for example “<unk>”

  • The output is a dictionary that you can use to convert each token into token ids

• Decode

  • This is pretty straightforward this step is usually done at the end. This converts the token id back into the word to output.

• Forward Pass

  • The forward pass takes the token ids seen so far and produces logits

  • Basically you take all the input tokens put them in the array and give it a score based on the logits that the model has prepared. This is the part we assume for the purposes of this article is a black box (Part of the model generating these logits using attention mechansims)

  • In the forward pass you create an array of the predicted words with its respective logits, these logits are raw scores for every possible next token in the vocabulary.

• Softmax and Temperature

  • Once you have the set of logits they are not a set of probabilites so we need to be able to convert them into probabilities for which you need the softmax function.

  • You can read about softmax here. https://en.wikipedia.org/wiki/Softmax_function

  • These set of probabilities are now easier to reason about and predict the next token against.

  • We have another concept here called as temperatures temperature controls how random the next-token selection is. We apply it by dividing each logit by a value T before softmax.

    • If T < 1, the distribution becomes sharper. The highest-scoring tokens become even more likely.

    • If T > 1, the distribution becomes flatter. Lower-scoring tokens get more chance, which can make the output more varied or creative.

• Sampling:

  • Once we have probabilities, we sample the next token. We pick a random number between 0 and 1, then use the cumulative probability ranges to decide which token gets selected.

  • For example, if the probabilities are [0.1, 0.2, 0.7], the ranges are:

    0.0 to 0.1 -> token 0

    0.1 to 0.3 -> token 1

    0.3 to 1.0 -> token 2

    A token with a larger probability owns a larger part of the 0 to 1 range, so the

    random number is more likely to land on it.

This is the basic loop of inference, i have a small github repo where i will upload the code for this, but this is the gist of it. In the future lets keep building on this base.