← Ashish's notebook/How Top-P Sampling Works in LLMs

Notebook

How Top-P Sampling Works in LLMs

Top-K keeps a fixed number of words every time. Top-P keeps fewer when the model is sure and more when it isn't.

18 SEP 2026

In the last notebook we learned about Top-K.

Top-K keeps a fixed number of the highest probability tokens and removes the rest.

Now we will learn Top-P (also called nucleus sampling). What is Top-P? Top-P is also a sampling technique of LLM which is done on the probabilities of the next token. But instead of keeping a fixed number of tokens, it keeps as many tokens as needed to reach a certain probability mass. So when the model is very confident it keeps fewer tokens, and when it is unsure it keeps more.

You will understand better by seeing it, so let's get right into it.

We will take the model apart again and actually see this happen.

Input → Tokenization → Model → Logits → Softmax → Probabilities → Sort → Cumulative sum → Top-P filtering → Sampling → Next token

What we will learn

1. Where Top-P Fits in Text Generation
2. What P means
3. How to Get the Model's Raw Logits
4. How Softmax Converts Logits into Probabilities
5. How to Sort the Probabilities
6. How the Cumulative Sum Decides the Cutoff
7. How to See the Nucleus (the kept tokens)
8. How Renormalization Changes the Probabilities
9. How Different P Values Change the Output
10. How Different P Values Change the Generated Text

Setup — Load the Model and Tokenizer

Same two steps as before: the tokenizer, which turns text into numbers, and the model itself.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen3-0.6B"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

Loading weights:   0%|          | 0/311 [00:00<?, ?it/s]

Let's take top_p=0.9 for now

from transformers import GenerationConfig

model.generation_config = GenerationConfig(
    do_sample=True,
    temperature=0.3,
    top_p=0.9,
    max_new_tokens=10,
)

Now we write the prompt and print the output of next predicted words.

return_tensors="pt" tells the tokenizer to return PyTorch tensors (vectors) instead of Python lists.

There are some special tokens in llms. For example - "<|endoftext|>", it tells the llm the end of text. We don't want these in the output , that's why we use skip_special_tokens=True it skips all the special tokens from the output.

prompt = "Can you bring me"

inputs = tokenizer(prompt,return_tensors="pt").to(model.device)

output = model.generate(**inputs)

print(tokenizer.decode(output[0],skip_special_tokens=True))

Can you bring me a list of all the possible combinations of the numbers

Chapter 1 — Get the logits

Before Top-P can happen, the model first produces logits.

The logits contain one score for every token in the vocabulary.

import torch

with torch.no_grad():
    outputs = model(**inputs)

logits = outputs.logits[0, -1, :]

print(logits.shape)

torch.Size([151936])

Chapter 2 — Convert logits to probabilities

Softmax converts the logits into probabilities.

Now every token has a probability between 0 and 1, and all probabilities add up to 1.

probs = torch.softmax(logits, dim=-1)

print(probs.shape)

torch.Size([151936])

Chapter 3 — Sort the probabilities

Top-P works on the tokens in order, from the highest probability to the lowest.

So first we sort all the probabilities in descending order. torch.sort gives us the sorted probabilities and also the original index of each token.

sorted_probs, sorted_idx = torch.sort(
    probs,
    descending=True
)

for prob, idx in zip(sorted_probs[:10], sorted_idx[:10]):
    token = tokenizer.decode([idx.item()]).strip() or "␣"
    print(f"{token!r:20} {prob.item():.4f}")

'a'                  0.1895
'the'                0.1670
'up'                 0.1152
'some'               0.1016
'help'               0.0422
'back'               0.0352
'an'                 0.0310
'your'               0.0273
'information'        0.0212
'to'                 0.0166

Chapter 4 — Cumulative probabilities

This is the main idea of Top-P.

We go down the sorted list and keep adding the probabilities one by one. This running total is called the cumulative probability.

torch.cumsum does this for us. The moment the cumulative sum crosses our P value, we stop. Every token up to that point is kept, and the rest are removed.

cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

for i in range(10):
    token = tokenizer.decode([sorted_idx[i].item()]).strip() or "␣"
    print(
        f"{token!r:20} "
        f"prob={sorted_probs[i].item():.4f}   "
        f"cumulative={cumulative_probs[i].item():.4f}"
    )
    

'a'                  prob=0.1895   cumulative=0.1895
'the'                prob=0.1670   cumulative=0.3555
'up'                 prob=0.1152   cumulative=0.4727
'some'               prob=0.1016   cumulative=0.5742
'help'               prob=0.0422   cumulative=0.6172
'back'               prob=0.0352   cumulative=0.6523
'an'                 prob=0.0310   cumulative=0.6797
'your'               prob=0.0273   cumulative=0.7109
'information'        prob=0.0212   cumulative=0.7305
'to'                 prob=0.0166   cumulative=0.7461

Chapter 5 — Pick the nucleus

The group of tokens we keep is called the nucleus. That's why Top-P is also called nucleus sampling.

Let's take top_p=0.5. We use torch.searchsorted to find the first place where the cumulative sum reaches 0.5, and we keep every token up to and including that one.

Lets see how many tokens it keeps. With top_p=0.5.

top_p = 0.5

cutoff = torch.searchsorted(
    cumulative_probs,
    torch.tensor(top_p, device=cumulative_probs.device)
).item()

top_p_probs = sorted_probs[:cutoff + 1]
top_p_idx = sorted_idx[:cutoff + 1]

print("Top-P:", top_p)
print("Tokens kept:", len(top_p_probs))
print("Probability:", top_p_probs.sum().item())

Top-P: 0.5
Tokens kept: 4
Probability: 0.57421875

And these are the tokens that made it into the nucleus:

tokens = [
    tokenizer.decode([idx.item()]).strip() or "␣"
    for idx in top_p_idx
]

for token, prob in zip(tokens, top_p_probs):
    print(f"{token!r:20} {prob.item():.4f}")

'a'                  0.1895
'the'                0.1670
'up'                 0.1152
'some'               0.1016

Chapter 6 — Visualize the nucleus

This is the graph we want to see.

We plot the top tokens and color them. The green bars are the tokens Top-P keeps (the nucleus for top_p=0.5) and the gray bars are the ones it removes.

import matplotlib.pyplot as plt

top_p = 0.5
n_show = 10

cutoff = torch.searchsorted(
    cumulative_probs,
    torch.tensor(top_p, device=cumulative_probs.device)
).item()

kept = cutoff + 1

show_probs = sorted_probs[:n_show].float().detach().cpu().numpy()
show_tokens = [
    tokenizer.decode([sorted_idx[i].item()]).strip() or "␣"
    for i in range(n_show)
]

colors = ["tab:green" if i < kept else "tab:gray" for i in range(n_show)]

plt.figure(figsize=(10, 5))
plt.bar(show_tokens, show_probs, color=colors)

plt.xlabel("Candidate next token")
plt.ylabel("Probability")
plt.title(f"Top-P = {top_p}  (green = nucleus, {kept} tokens kept)")
plt.xticks(rotation=45)

plt.show()

Chapter 7 — Original vs Top-P

This part is important. We will learn about renormalization which is an important part.

After Top-P keeps the nucleus, the probabilities of those tokens no longer add up to 1. For top_p=0.5 our kept tokens only add up to around 0.57.

For sampling, the remaining probabilities are renormalized so that they add up to 1. Each kept probability is divided by the sum of the kept probabilities:

0.1895 / 0.5742
0.1670 / 0.5742
0.1152 / 0.5742
0.1016 / 0.5742

So:

Top-P picks the tokens. Renormalization gives those tokens their new probabilities.

renormalized_probs = top_p_probs / top_p_probs.sum()

for token, original, new in zip(
    tokens,
    top_p_probs,
    renormalized_probs
):
    print(
        f"{token!r:20} "
        f"original={original.item():.4f}   "
        f"after Top-P={new.item():.4f}"
    )

print("Sum after renormalization:", renormalized_probs.sum().item())

'a'                  original=0.1895   after Top-P=0.3301
'the'                original=0.1670   after Top-P=0.2910
'up'                 original=0.1152   after Top-P=0.2002
'some'               original=0.1016   after Top-P=0.1768
Sum after renormalization: 1.0

Chapter 8 — Compare different values of P

Now let's actually see what changes when we use:

P = 0.5
P = 0.7
P = 0.9
P = 0.95
P = 0.99

Remember:

• small P → fewer tokens kept (safer output)
• large P → more tokens kept (more variety)

Unlike K, the number of tokens is not fixed. It depends on how confident the model is at that step.

for p in [0.5, 0.7, 0.9, 0.95, 0.99]:

    cutoff = torch.searchsorted(
        cumulative_probs,
        torch.tensor(p, device=cumulative_probs.device)
    ).item()

    print(f"top_p={p}: {cutoff + 1} tokens kept")

top_p=0.5: 4 tokens kept
top_p=0.7: 8 tokens kept
top_p=0.9: 109 tokens kept
top_p=0.95: 704 tokens kept
top_p=0.99: 6274 tokens kept

Chapter 9 — Output generation with different values of P

Here we are setting different values of P and seeing how the output changes based on that

for p in [0.5, 0.9, 0.95]:

    model.generation_config = GenerationConfig(
        do_sample=True,
        temperature=0.3,
        top_p=p,
        max_new_tokens=10,
    )

    output = model.generate(**inputs)

    text = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    print(f"Top-P={p}: {text}")

Top-P=0.5: Can you bring me a list of all the possible combinations of 3
Top-P=0.9: Can you bring me the list of all the possible combinations of the numbers
Top-P=0.95: Can you bring me some information about the history of the 19

Summary

Top-P (nucleus sampling) is a sampling method that keeps the smallest group of top tokens whose probabilities add up to a chosen value P, and removes the rest. We went through the process step by step: the model first produces logits, softmax converts them into probabilities, we sort them, and torch.cumsum gives the cumulative sum. Top-P keeps every token until the cumulative sum reaches P. Since the kept probabilities no longer add up to 1, they are renormalized so the model can sample from them properly.

We also saw how different P values affect the output. A small P, such as P=0.5, keeps very few tokens and makes the output more predictable, while a larger P keeps more tokens and gives more variety. The key difference from Top-K is that the number of tokens is not fixed. It grows or shrinks depending on how confident the model is at each step.

← All notebooks

GitHub LinkedIn