Building a Mini Inference Engine from Scratch

Most of us know that LLMs are auto-regressive, i.e., they generate token one by one. We also know that they have a concept of KV caching and can process multiple requests together, requests of different sizes, producing different amounts of output tokens. Reading these concepts feels good & easy because we can just say they work like this or they work that, but how does it actually happen?

• How do they produce tokens one by one?
• How do they decide that this is the next token?
• How does KVCache actually get implemented?
• How do to process batching?

These are some questions which I have had in mind for a long time. And I wanted to realize them by implementing them myself. So in this artifact, we are going to build a mini inference system from scratch, which will help us answer these questions from the first principles :)

So lets get started, but before that, I already created a proper YouTube Series on building the LLM inference engine from scratch. So, if you want a visual explanation and if you understand the concepts better by videos, I hope going through this series will be much helpful for you to visualise and understand the whole implementation. Or you may follow your own custom combination, its all upto you.

Ok, let us first understand what actually happens when we provide the prompt or input tags to the model.So first step, that your token goes through: you gave a prompt to ChatGPT, Grok, Cloude, or whatever LLM of your choice, assuming that it is an auto-regressive LLM. The first step it goes through is the tokenizer, which converts input text into token IDs and matrices. It cannot process text like us, so it needs the input to be present in the form of numbers. So to make it understandable for the LLM, the input gets converted into token IDs. If you want to know more or understand, you may go to the OpenAI tokenizer to get a real feel of how it works in GPT.

And token ID is basically referred to the tokens present in the vocabulary set of the model. So essentially, we have converted the text into tokens which the model can understand. After this step, we convert the tag tokens to vector embeddings. What does that mean? Vector is basically a mathematical entity having features. And in case of a model, it is simply converting into a row all columns of numbers.

And this collection of all the input vectors is something we represent by X.