TL;DR — This post walks you through building a temperature‑aware top‑p nucleus sampler with top‑k fallback in pure Python, giving you a concrete portfolio piece that demonstrates algorithmic understanding, performance tuning, and production‑ready engineering. You will end up with a runnable module that can be integrated into larger LLM inference pipelines and easily extended with persistence, observability, and scaling.

The rise of large language models has turned token sampling from a theoretical curiosity into a critical engineering challenge. Whether you are serving a model on a GPU‑cluster or embedding it in a mobile app, the choice of sampling strategy directly affects latency, diversity, and user perception. In this guide you will implement a temperature‑aware top‑p nucleus sampler with a top‑k fallback from scratch, using only the Python standard library. The resulting module is small enough to fit in a résumé bullet, yet powerful enough to be dropped into a real inference service.

Why This Project Stands Out on a CV

  • Algorithmic depth – You will write the core probability transformations (softmax, temperature scaling, cumulative distribution) yourself, showing you understand the math behind modern LLM decoding.
  • Systems thinking – The sampler is designed as a composable function that can be plugged into a larger inference pipeline, mirroring how production services like vLLM or TensorRT‑LLM structure their sampling kernels.
  • Performance awareness – By operating on NumPy‑free plain Python lists and leveraging built‑in heapq for top‑k selection, you demonstrate an ability to reason about constant factors and memory usage—skills that surface in system design interviews.
  • Fallback strategy – Implementing a top‑k fallback when the nucleus threshold is too aggressive highlights fault‑tolerant design, a common requirement in high‑availability serving systems.
  • Portfolio signal – A GitHub repo containing this sampler, accompanied by benchmarks against Hugging Face’s generate API, provides concrete evidence of end‑to‑end engineering ability.

Architecture Overview

The sampler is composed of three logical layers:

  1. Logit preprocessing – Receives raw model logits (a list of floats) and applies temperature scaling.
  2. Selection core – Implements two complementary strategies:
    • Top‑k truncation – Keeps only the k highest‑probability tokens.
    • Nucleus (top‑p) truncation – Removes the least probable tokens until the cumulative probability exceeds p.
  3. Fallback orchestrator – Combines the two strategies: if the nucleus set is empty (e.g., p too low), it falls back to the top‑k set, guaranteeing a non‑empty candidate list.

A simple text diagram:

Logits ──► Temperature Scaling ──► Top‑k Filter ──► Candidate Set
                     │
                     └──► Nucleus Filter ──► Candidate Set (if non‑empty)

The final step draws a token from the candidate set using inverse transform sampling, returning both the token id and its probability.

Building It Step by Step

Below are the numbered steps to implement the sampler. Each snippet is self‑contained and can be run in a Python 3.10+ environment.

Step 1: Temperature Scaling

import math
from typing import List

def softmax(logits: List[float], temperature: float) -> List[float]:
    """Compute softmax probabilities with temperature scaling."""
    if temperature <= 0:
        raise ValueError("Temperature must be positive.")
    # Subtract max for numerical stability
    max_logit = max(logits)
    exps = [math.exp((logit - max_logit) / temperature) for logit in logits]
    sum_exps = sum(exps)
    return [exp / sum_exps for exp in exps]

Step 2: Top‑k Selection

import heapq

def top_k_indices(probs: List[float], k: int) -> List[int]:
    """Return indices of the k largest probabilities."""
    if k <= 0:
        raise ValueError("k must be positive.")
    # Use a min‑heap to keep the k largest elements
    heap = []
    for idx, prob in enumerate(probs):
        if len(heap) < k:
            heapq.heappush(heap, (prob, idx))
        else:
            if prob > heap[0][0]:
                heapq.heapreplace(heap, (prob, idx))
    # Extract indices sorted by descending probability
    return [idx for _, idx in sorted(heap, key=lambda x: -x[0])]

Step 3: Nucleus (Top‑p) Selection

def top_p_indices(probs: List[float], p: float) -> List[int]:
    """Return indices forming the smallest set whose cumulative probability >= p."""
    if not (0 < p <= 1):
        raise ValueError("p must be in (0, 1].")
    # Sort indices by descending probability
    sorted_indices = sorted(range(len(probs)), key=lambda i: -probs[i])
    cumulative = 0.0
    selected = []
    for idx in sorted_indices:
        cumulative += probs[idx]
        selected.append(idx)
        if cumulative >= p:
            break
    return selected

Step 4: Fallback Orchestrator

def sample_token(
    logits: List[float],
    temperature: float = 1.0,
    top_k: int = 50,
    top_p: float = 0.9,
) -> int:
    """
    Sample a token using temperature‑aware top‑p nucleus sampling
    with top‑k fallback.
    Returns the token index.
    """
    probs = softmax(logits, temperature)
    # Try nucleus first
    nucleus = top_p_indices(probs, top_p)
    if nucleus:
        candidates = nucleus
    else:
        # Fallback to top‑k
        candidates = top_k_indices(probs, top_k)
    # Inverse transform sampling from the candidate set
    total = sum(probs[i] for i in candidates)
    r = random.random() * total
    cumulative = 0.0
    for i in candidates:
        cumulative += probs[i]
        if r <= cumulative:
            return i
    # Should never reach here, but return last candidate as safety
    return candidates[-1]

Step 5: Putting It All Together in a Module

Create a file nucleus_sampler.py:

# nucleus_sampler.py
import heapq
import math
import random
from typing import List

def softmax(logits: List[float], temperature: float) -> List[float]:
    # ... (implementation from Step 1)

def top_k_indices(probs: List[float], k: int) -> List[int]:
    # ... (implementation from Step 2)

def top_p_indices(probs: List[float], p: float) -> List[int]:
    # ... (implementation from Step 3)

def sample_token(
    logits: List[float],
    temperature: float = 1.0,
    top_k: int = 50,
    top_p: float = 0.9,
) -> int:
    # ... (implementation from Step 4)

Running and Testing It

You can verify the sampler with a quick script:

python -c "
from nucleus_sampler import sample_token
logits = [0.5, 1.2, 0.8, 2.0, 0.1]
print('Sampled token:', sample_token(logits, temperature=0.8, top_k=3, top_p=0.95))
"

For a more thorough test, generate a synthetic vocabulary of 10 000 tokens and compare the empirical distribution against the expected softmax when temperature=1.0, top_k=100, top_p=0.9. Use pytest with a fixture that seeds random for reproducibility.

# test_sampler.py
import random
import pytest
from nucleus_sampler import sample_token, softmax

def test_temperature_scaling():
    logits = [1.0, 2.0, 3.0]
    probs = softmax(logits, temperature=2.0)
    assert abs(sum(probs) - 1.0) < 1e-6

def test_fallback_when_nucleus_empty():
    # Force top_p extremely low so nucleus returns empty
    logits = [0.1, 0.2, 0.3]
    random.seed(42)
    token = sample_token(logits, temperature=1.0, top_k=2, top_p=0.01)
    assert token in (0, 1)  # top‑2 indices

Run the tests with:

pytest -q test_sampler.py

All tests should pass, confirming that the sampler behaves correctly across edge cases.

Extending It: Your Roadmap to Senior-Level

  1. Persisted sampling state – Store temperature, top‑k, and top‑p in a configuration file (e.g., YAML) and load it at startup; this matters because production services often need to A/B‑test decoding hyperparameters without code changes.
  2. Horizontal scaling – Wrap the sampler in a FastAPI endpoint and deploy multiple replicas behind a load balancer; each replica can handle a request independently, enabling throughput scaling.
  3. Observability – Emit Prometheus metrics for sampling latency, fallback frequency, and token probability distributions; these metrics are essential for detecting drift in model behavior.
  4. Fault tolerance – Introduce a retry loop with exponential backoff when the underlying model service is unavailable, and circuit‑breaker logic to prevent cascading failures.
  5. Benchmarking harness – Create a script that measures tokens‑per‑second and perplexity across a range of temperature and top‑p values, using a fixed prompt and comparing against Hugging Face’s generate as a baseline.
  6. Batched sampling – Extend sample_token to accept a 2‑D tensor of logits (batch × vocab) and return a vector of sampled token ids, leveraging vectorized operations for GPU‑accelerated inference.

Each upgrade not only deepens the project’s technical scope but also aligns it with the expectations of senior engineering roles: reliability, scalability, and data‑driven decision making.

Key Takeaways

  • Implementing the core probability transformations yourself demonstrates algorithmic mastery.
  • Designing a fallback strategy showcases fault‑tolerant thinking.
  • The module is lightweight, making it easy to integrate into larger inference pipelines.
  • Extending with persistence, scaling, observability, fault tolerance, and benchmarking prepares the project for production use.
  • A well‑documented GitHub repo with tests and benchmarks provides concrete evidence of systems engineering capability.

Further Reading