Skip to content

Models

Almost all of the LLM models today are autoregressive models that do the task of next token prediction. They operate over sequences of tokens, and the fewer the tokens to generate, the faster is the end to end inference.

Working with Sequences

Inference involves two or three types of sequences of tokens

  • Input: Prompt, chat, context, function calls and any other type of input passed to the LLM
  • Reasoning: Optionally, for reasoning models, an intermediate output sequence for thinking
  • Output: The response generated by LLM

Combined, these sequences together are limited by the model's context window. We can additionally have a max tokens argument to denote the length of output after which generation ceases.

While input sequence is a string, LLMs are trained to accept a variety of inputs like multi turn chat sequences with roles, function signatures for tool calls, and in some cases multimodal inputs like images. This is all handled by a chat template and differs subtly from model to model, and must be implemented correctly in the inference engine.

Tokenizing input sequence, with chate template applied, step zero of inference. Then, there are two primary phases of inference

  • Prefill: Process the input sequence to calculate attention for each input token and store these values in a KV cache
  • Decode: Perform forward pass through the model to generate tokens autoregressively

There are a couple of ways to influence which token is picked: Temperature, Top-K, Top-p

Additional tools like logit biasing are also used for further directing the output of an LLM, especially for generating structured outputs like JSON for performing tool calls.

LLM Model Variants

There are several different architectures of models implemented by various companies and open source communities. For a given architecture type, there are several variants of models avaialble (especially on platforms like huggingface).

  • Number of Parmeters: These vary based on number of model parameters (8B vs 70B variants)
  • Model Type: Base vs Instruct variants
  • Quantization: Floating point precision vs integer
  • Fine Tuned Variants: For tasks in specialized domains, using methods like PEFT, LoRA

Base vs Instruct

The base model is primarily trained for the task of next token prediction. On the other hand, the instruct model is additionally trained/aligned to follow human instructions and produce useful conversational responses. Both the models typically have the same underlying capability, but harness it in different ways. Instruct model is usually more optimized for chat responses and more reliable in structured output generation.

Transformer Blocks

The core transformer blocks are

  • Embedding Layer: The input layer that takes in the sequence of tokens, and produces a sequence of embedding vectors
  • Transformer Blocks: The main component that implements attention mechanisms, the aim of these blocks is to transform the vectors from the embedding space to a different space/representation with more semantic context infused
  • Output Layer: Generates the logits corresponding to the tokens in vocabulary, based on which the next token is selected

Within a transfomer block als we have attention, normalization and feed forward network layers. The feed forward network layer makes the majority of the parameters of the network, followed by the attention blocks.

Attention

\[ \begin{aligned} Attention(K,Q,V) = Softmax \left(\frac{QK^{T}}{\sqrt{D}}\right)V \end{aligned} \]

Types of attention

  • Self Attention: Key, Query and Value vectors are from the same sequence
  • Masked Self Attention: Mainly used in decoder only transformers, wherein, a token attends to only the previous tokens
  • Cross Attention: Mainly used in encoder-decoder style of architectures, where Key and Value vectors come from the input sequence, and the Query vectors come from the output sequence

Attention has a quadratic dependence on the input sequence. If we use KV cache, the attention operation becomes a simple lookup and the operation becomes linear, allowing better scaling of transformers on longer sequences.

KV Cache is built during the prefill phase, and used and updated during the decode phase. The cache lives in the GPU memory by default.

Mixture of Experts (MoE)

Density of a network depends on the number of connections between different layers. Dense networks can store more information, but sparse networks take less compute and memory to run.

Mixture of Experts (MoE) is an architecture optimization that uses hundreds of smaller matrices (experts) instead of one giant matrix. It routes each input to a small selection of experts and this process is called activation of experts.

Consider the model Qwen3-235B-A22B. Here the last part of the name means that 22B parameters out of 235B parameters are activated per request.

However, in batched requests on servers, different requests might activate different experts depending on the input. This means almost all of the model parameters might be active, unless sparsity is achieved by large scale expert parallelism.

Expert routing is granular. The router, a tiny model within the LLM picks which experts to activate at each layer of the model. In the Quen example, there are 128 experts avaialble, and router picks eight experts at each of the 94 layers for every token that is generated.

MoE is popular for models with 100B+ parameters. MoE unlocks a new form inference parallelism called Expert Parallelism which enables high throughput inference for large models on multiple GPUs.

Models under 8B parameters tend to use traditional dense architecture efficiently. Domain specific models dont gain much from MoE as the entire model is one expert.