My LLM Generates Text 🚀 Now Let’s Make It Faster 🚀

In my last post, we saw how an LLM generates text one token at a time:

1. Read the prompt
2. Score every possible next token
3. Pick one token
4. Add it to the prompt
5. 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:

1. Let PyTorch use the best device available
2. Try torch.compile() and understand the warm-up run
3. 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.