In my last post, we saw how an LLM generates text one token at a time:
- Read the prompt
- Score every possible next token
- Pick one token
- Add it to the prompt
- Repeat
That is cool, but there is a small problem: it can be slow.
If you did not read the first post, start here: How an LLM Generates Text, One Token at a Time.
Today we are going to make the same small Qwen model faster in three steps:
- Let PyTorch use the best device available
- Try
torch.compile()and understand the warm-up run - Use a KV cache so the model does not process the old prompt again and again
No training. Same model. Same greedy token selection. We are only changing how efficiently it runs.
Pre-requisites
This uses the same setup as the first article:
pip install -U torch transformers The code downloads the open Qwen/Qwen3-0.6B-Base model the first time you run it. CPU works fine for learning, but a GPU makes the difference easier to see.
Step 1: let the code pick the device
Before loading the model, we ask PyTorch what hardware we can use.
def get_device(): if torch.cuda.is_available(): print("Using NVIDIA CUDA GPU") return torch.device("cuda") if torch.backends.mps.is_available(): print("Using Apple Silicon GPU (MPS)") return torch.device("mps") print("Using CPU") return torch.device("cpu") Easy peazy:
- NVIDIA GPU? Use
cuda. - Mac with Apple Silicon? Use
mps. - Nothing else? CPU still works.
Then we move both the model and input tokens to that device:
device = get_device() model.to(device) input_token_ids = torch.tensor( tokenizer.encode(prompt), device=device ).unsqueeze(0) The model and its tensors need to live in the same place. A model on the GPU cannot process tokens that are still sitting on the CPU.
Step 2: first, our original generation loop
This is the same simple loop from the last article. It sends every token written so far back through the model on every iteration.
@torch.inference_mode() def generate_text_basic(model, token_ids, max_new_tokens, eos_token_id): input_length = token_ids.shape[1] model.eval() for _ in range(max_new_tokens): logits = model(input_ids=token_ids).logits[:, -1] next_token = torch.argmax(logits, dim=-1, keepdim=True) if torch.all(next_token == eos_token_id): break token_ids = torch.cat([token_ids, next_token], dim=1) return token_ids[:, input_length:] Imagine the model has generated three tokens already:
Prompt + token 1 + token 2 + token 3 To generate token 4, this basic loop reads all of that again. It works, but it repeats work.
Step 3: what happens when we use torch.compile()?
torch.compile() asks PyTorch to look at the model operations and try to run them more efficiently. It does not train the model and it does not change its answers.
model_compiled = torch.compile(model, dynamic=True) I am using dynamic=True because our input becomes longer as new tokens are added. Without it, the compiler may try to recompile when it sees a new sequence length.
Now let’s run the same prompt three times:
for run in range(3): start_time = time.time() output_token_ids = generate_text_basic( model_compiled, input_token_ids, max_new_tokens=100, eos_token_id=tokenizer.eos_token_id, ) total_time = time.time() - start_time if run == 0: print("Warm-up run finished") else: tokens_per_second = output_token_ids.numel() / total_time print(f"Run {run}: {tokens_per_second:.1f} tokens/second") Why ignore the first result? The first run has extra work: PyTorch is preparing the compiled version of the model. The second and third runs are the ones worth comparing.
This does not remove the repeated prompt processing. It only tries to make those calculations faster.
Step 4: now remove compile and use KV cache
For this next test, use the original model, not model_compiled:
# This is the normal, non-compiled model. model.eval() Then use Hugging Face generation with caching turned on:
@torch.inference_mode() def generate_text_with_cache(model, token_ids, max_new_tokens, eos_token_id): all_token_ids = model.generate( input_ids=token_ids, max_new_tokens=max_new_tokens, do_sample=False, use_cache=True, eos_token_id=eos_token_id, ) # `generate` returns the prompt plus new tokens. return all_token_ids[:, token_ids.shape[1]:] The result is still greedy decoding:
pick the highest-scoring token → append it → repeat But use_cache=True tells the model to save its previous attention work.
Without cache:
Prompt → predict token 1 Prompt + token 1 → predict token 2 Prompt + token 1 + token 2 → predict token 3 With cache:
Prompt → process once and save attention work token 1 → use saved work token 2 → use saved work token 3 → use saved work The model still has access to the full prompt. The cache stores the keys and values it needs from the old tokens, so it does not need to calculate them again.
Run it like this:
start_time = time.time() output_token_ids = generate_text_with_cache( model, input_token_ids, max_new_tokens=100, eos_token_id=tokenizer.eos_token_id, ) total_time = time.time() - start_time tokens_per_second = output_token_ids.numel() / total_time print(f"KV cache: {tokens_per_second:.1f} tokens/second") print(tokenizer.decode(output_token_ids[0], skip_special_tokens=True)) So which one is better?
They improve different things:
| Method | What it changes |
|---|---|
| GPU / MPS | Runs the same calculations on faster hardware |
torch.compile() | Tries to execute the same model operations more efficiently |
| KV cache | Stops recalculating attention data for old tokens |
KV cache is usually the bigger win for token-by-token generation, because it removes repeated work instead of only making that repeated work faster.
Your numbers will depend on your model, prompt length, PyTorch version, and hardware. Benchmark your own machine instead of trusting my numbers.
The complete example
import time import torch from transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = "Qwen/Qwen3-0.6B-Base" def get_device(): if torch.cuda.is_available(): print("Using NVIDIA CUDA GPU") return torch.device("cuda") if torch.backends.mps.is_available(): print("Using Apple Silicon GPU (MPS)") return torch.device("mps") print("Using CPU") return torch.device("cpu") @torch.inference_mode() def generate_text_basic(model, token_ids, max_new_tokens, eos_token_id): input_length = token_ids.shape[1] model.eval() for _ in range(max_new_tokens): logits = model(input_ids=token_ids).logits[:, -1] next_token = torch.argmax(logits, dim=-1, keepdim=True) if torch.all(next_token == eos_token_id): break token_ids = torch.cat([token_ids, next_token], dim=1) return token_ids[:, input_length:] @torch.inference_mode() def generate_text_with_cache(model, token_ids, max_new_tokens, eos_token_id): all_token_ids = model.generate( input_ids=token_ids, max_new_tokens=max_new_tokens, do_sample=False, use_cache=True, eos_token_id=eos_token_id, ) return all_token_ids[:, token_ids.shape[1]:] def print_stats(name, output_token_ids, start_time, tokenizer): total_time = time.time() - start_time tokens_per_second = output_token_ids.numel() / total_time text = tokenizer.decode(output_token_ids[0], skip_special_tokens=True) print(f"{name}: {tokens_per_second:.1f} tokens/second") print(text) device = get_device() tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to(device) model.eval() prompt = "Explain large language models in a single sentence." input_token_ids = torch.tensor( tokenizer.encode(prompt), device=device ).unsqueeze(0) max_new_tokens = 100 # Test 1: compile the original loop three times. model_compiled = torch.compile(model, dynamic=True) for run in range(3): start_time = time.time() output_token_ids = generate_text_basic( model_compiled, input_token_ids, max_new_tokens, tokenizer.eos_token_id, ) if run == 0: print("Compiled warm-up run finished") else: print_stats(f"Compiled run {run}", output_token_ids, start_time, tokenizer) # Test 2: remove compile and let `generate` use its KV cache. start_time = time.time() output_token_ids = generate_text_with_cache( model, input_token_ids, max_new_tokens, tokenizer.eos_token_id, ) print_stats("KV cache", output_token_ids, start_time, tokenizer) One last thing
This article uses a normal growing KV cache because it is the easiest version to understand. A growing cache is great for generation, but it can make torch.compile() harder because the cache shape changes every token.
That is not a problem we need to solve today.
Today, the important part is this:
GPU makes calculations faster. Compile tries to make operations faster. KV cache avoids doing old calculations again. And that is how we start making an LLM faster.
For more detail about cache strategies and their memory/speed trade-offs, see the official Hugging Face KV-cache documentation.
