Building a Qwen2.5-1.5B Inference Engine From Scratch
This article walks through the implementation of an inference engine for Qwen2.5-1.5B from the ground up. Along the way, I'll unpack transformer architecture, how an inference engine differs from the model itself, and why anyone would bother implementing one from scratch when highly optimized alternatives already exist.
This is also my first technical article since my peak in high school (hopefully it's not too evident). This process has allowed me to confront gaps in my own understanding that I hope to make accessible to others along the way. By the end of this article, I will showcase a more concrete intuition for the systems that make inference efficient. Rather than treating each optimization as a black box, I'll introduce them incrementally, benchmarking each change to understand not just whether it improves performance, but why and under what conditions.
1. Introduction
I'm under the assumption that most, if not all, of you reading this article are familiar with the great Attention Is All You Need paper. Or developed sufficient knowledge by participating in enough conversations where everyone else has also allegedly read it. If not, I'll cover the relevant components as we implement them.
Inference is the process of running a trained model with learned weights on new inputs, to produce new outputs. Bottlenecks arise when factors like computational demands and memory bandwidth limit how efficiently and effectively these predictions can be generated.
While a model defines the computations needed to produce these predictions, it doesn't specify how to execute them efficiently on various types of hardware and meet the performance requirements of a given use case. An inference engine is responsible for identifying how to execute these operations efficiently, managing things like weights and intermediate states in memory, reducing redundant computations, and orchestrating incoming requests.
Most people in this space separate inference into two stages – prefill and decode. During prefill, the model processes the input prompt, computes its representations, and populates the KV cache. Decode begins when the model starts generating new tokens for the new input using the cached keys and values.
I think an important distinction that is important to understand is the difference between autoregressive and sequential computation. Autoregression means the model generates each new token from its own previous outputs, feeding them back as input. Sequential is an execution constraint: one computation must finish before the next can begin. Autoregressive generation creates a sequential dependency across output tokens, but not all computation in an autoregressive model is sequential.
During prefill, all input tokens are already known, so the model can process them in parallel and makes prefill largely compute-bound. During decode, each new token must exist before the next can be predicted, so tokens are generated one at a time which makes decode largely memory-bandwidth-bound. The point being: because the two phases have different bottlenecks, an optimization that improves prefill performance may have little effect on decode, and vice versa.
The goal of this project is to build an inference engine for Qwen2.5-1.5B-Instruct from scratch and make it measurably faster than a naive implementation of the same model. Given the different bottlenecks in prefill and decode, we'll explore optimizations targeting each stage and measure their impact on time to first token (TTFT), decode throughput (tokens per second), and peak GPU memory usage.
1.1 Loading the model and establishing a reference
I
downloadedfrom transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float32
).to("cuda")
the pretrained Qwen2.5-1.5B-Instruct weights and tokenizer from
Hugging Face and loaded the model in FP32.
Google Colab’s NVIDIA T4 GPUs don't support native BF16 arithmetic, so I loaded the model in FP32. The increased precision of 32 bits introduces less rounding error during calculations and make it easier to compare results. There is one caveat that FP32 requires twice as much memory as BF16 or FP16, which is why I’ll later experiment switching to FP16 and its impact on performance.
Qwen2.5-1.5B-Instruct contains 1,543,714,304 parameters distributed across 28 transformer layers, with 12 query attention heads and 2 key-value heads. I will be explaining these concepts in the following sections.
To establish a reference, I
ranimport json, torch
prompts = [
"Is Stanford a real school?",
"Explain what a transformer is in simple words.",
"User: Hi, how are you?\nAssistant: I'm doing well, thank you!\nUser: Can you tell me a joke about inference? Not that inference is a joke.",
"Write a Python function that adds two numbers.",
"¿Por qué hay tantas ardillas en Berkeley?",
"If a student has 3 midterms and 2 more get scheduled for the same week, how many does she have?",
"Write a very short poem about Dubai chocolate.",
"List three benefits of studying in the morning. Keep it short.",
"What is the difference between a list and a tuple in Python? Explain it like a TA at office hours with 40 people in the queue.",
"Why is the sky blue?",
]
model.eval()
ground_truth = []
for prompt in prompts:
text = tokenizer.apply_chat_template([{"role": "user", "content": prompt}],
tokenize=False, add_generation_prompt=True)
input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device)
with torch.no_grad():
out = model.generate(input_ids, max_new_tokens=100, do_sample=False, use_cache=False)
logits = model(input_ids).logits[0, -1]
gen_ids = out[0, input_ids.shape[1]:].tolist()
ground_truth.append({
"prompt": prompt,
"input_ids": input_ids[0].tolist(),
"generated_ids": gen_ids,
"first_token_logits": logits.tolist(),
})
print(prompt, "→", tokenizer.decode(gen_ids, skip_special_tokens=True)[:80])
json.dump(ground_truth, open("ground_truth.json", "w"))
the downloaded model on ten prompts spanning different tasks.
Formatted with the given chat template, each prompt was tokenized
and saved along with the output, as well as the logits for the
first generated token. A naive comparison of the output tokens
could suffice, but a more robust comparison involves identifying if
the forward pass generated the same logit, the raw, unnormalized
score representing how strongly the model favors a particular token as
the next prediction, for every token in its vocabulary. This allows
us to check that both implementations assign nearly identical
scores to each possible next token rather than simply selecting the
same highest-scoring token.
Here's a saved example, formatted for readability.
2. Reimplementing Qwen's Transformer Architecture
Now that the reference outputs have been created, I'll reproduce
them by implementing the underlying operations using the
pretrained weightsweights = {}
for name, param in model.named_parameters():
weights[name] = param.detach().clone(). Qwen2.5-1.5B is a decoder-only transformer
model with 28 transformer layers, each with an attention layer and
a feed-forward network (see diagram below). I'll finally compare
the outputs of this reproduction against the reference and
identify opportunities for optimization.
2.1 Tokens to Embeddings
When a transformer receives an input, it needs to convert the text into a numerical representation that can be processed through a series of operations like matrix multiplication. This conversion process happens through tokenization and embedding. These processes are deterministic - there is no learning done here. Both phases done on the same input will always provide the same output.
Tokenization is when text is broken down into chunks that are then converted into unique integer IDs through a predefined dictionary. Qwen2.5 uses a byte-level Byte Pair Encoding (BPE) tokenizer.
To build the intuition behind this, consider the scenario where a word maps directly to a single ID. Because there is a practically unlimited number of possible words, including new words that users can create, maintaining a one-to-one mapping would require an impractically large vocabulary. While an ID mapped to a single character can seem like a more plausible approach, supporting different languages and new characters is also difficult.
Unicode assigns each character a number called a code point, which UTF-8 encodes into a sequence of one to four bytes. Code points above 127 are split into more than one byte. This allows us to represent any word without allocating a separate vocabulary entry to every character or word.
However, representing every character as individual bytes can introduce long sequence lengths which increases the amount of computation needed during inference. Byte Pair Encoding (BPE) resolves this by merging frequently occurring adjacent sequences of bytes into larger tokens. Here's an article that explains how the merge rules are created at a more granular level.
Since tokenization runs once per prompt on the CPU and takes a
negligible fraction of total inference time, I
reusedtext = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
tokenize=False, add_generation_prompt=True
)
input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device)
HF’s
tokenizer rather than reimplementing it, just as inference engines
like vLLM have done. Here’s a small interactive tokenizer
using the existing pretrained one you can experiment with below.
Loading Qwen’s vocabulary…
Once BPE has produced the final tokens, each token is mapped to its corresponding integer ID in Qwen's vocabulary table. These IDs are then used to retrieve the learned embedding vectors that the transformer block uses.