A reproducible sampler with top-k, top-p and speculative acceptance

Quillmate (fictional) is a fiction-writing app. When a user presses "regenerate" with the same seed, they must get the same text: after a server restart, and whichever other requests share the batch. You write the sampler. Python 3, NumPy and the standard library only. No AI tools.

Rules for one row of logits:

1. If temperature is 0, return the argmax (lowest id on ties) and ignore the other settings.
2. Divide the logits by the temperature.
3. If top_k > 0, keep the k highest logits. Order by logit (high first), then by token id (low first).
4. If top_p < 1, take the probabilities of the tokens kept so far, renormalized so they sum to 1. Order them by probability (high first), then by id (low first), and keep the smallest prefix whose total is at least top_p. Always keep at least one token.
5. Renormalize. Walk the kept tokens in ascending id order, adding up probability, and return the first token whose running total is above u, a uniform number in [0, 1).
6. Logits of minus infinity mark banned tokens. They are never returned.

Part 1 (15 minutes). sample_token(logits, temperature, top_k, top_p, u) -> int.

Part 2 (15 minutes). Reproducible streams.

• uniform(seed, step, stream=0) -> float: a uniform number in [0, 1) that depends only on its arguments.
• Generator(model, prompt, seed, temperature=1.0, top_k=0, top_p=1.0) with step() -> int, state() -> dict (JSON-safe) and Generator.from_state(model, state). model(tokens) returns the logits for the next token.
• A run that saves its state after 7 tokens and resumes must produce the same 15 tokens as a straight run.

Part 3 (15 minutes). sample_batch(logits, temperature, top_k, top_p, u) for a (B, V) array with settings per row, vectorized with NumPy (no Python loop over rows). It must return the same token as sample_token for every row. Then step_batch(generators): a request's tokens must not depend on which other requests share the batch, or in what order.

Part 4 (10 minutes). Speculative decoding with one draft token: speculative_step(p_target, q_draft, draft_token, u_accept, u_resample) -> (token, accepted). Accept the draft with probability min(1, p/q). If it is rejected, sample from max(0, p − q), renormalized, using u_resample. The output must follow p_target exactly.

Example: logits = log([0.05, 0.5, 0.15, 0.3]), temperature 1, top_k 0, top_p 0.75. By probability: id 1 (0.5), id 3 (0.3), id 2 (0.15), id 0 (0.05). 0.5 is below 0.75, and 0.5 + 0.3 = 0.8 reaches it, so keep ids 1 and 3, renormalized to 0.625 and 0.375. With u = 0.5, id 1's running total (0.625) is above 0.5, so return 1. With u = 0.7, return 3.

Write the order in which you apply temperature, top-k and top-p, and how you break ties. Then write the state a request needs so it can resume after a crash and produce the same tokens.