Fix frozen PRNG in categorical_sampling under repeated sampling

[utkarshtiwari-24]

Problem

With temperature > 0, repeated requests to mlx_lm.server return identical text. The sampling stops being random after the first request.

Root cause

categorical_sampling is decorated with @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state). In the server's generation loop the RNG state stops advancing after the first request, so every sample reuses the same state. Calling categorical_sampling or generate() directly does not hit this (the state advances normally), so it only appears in the long-running compiled server path.

Fix

Remove the decorator. The function is a single mx.random.categorical call, so the compile gives little benefit, and without it the function uses the live global RNG state on each call.

Repro

1. Start the server:
   mlx_lm.server --model mlx-community/Qwen2.5-0.5B-Instruct-4bit

2. Run the below curl cmd twice:
   curl -s localhost:8080/v1/chat/completions -H "Content-Type: application/json"
   -d '{"messages":[{"role":"user","content":"Tell me a one-sentence story."}],"temperature":1.0,"max_tokens":40}'

Before, two identical requests return byte-identical completions:

1. Once upon a time, in a neighboring village, a young woman named Rose found herself lost ...
2. Once upon a time, in a neighboring village, a young woman named Rose found herself lost ...

After, they differ:

1. In a land of empires ... a hero named Robin ... a mighty rebellion against the Elite Nobility ...
2. Ah, a sample story for you: When Eliza first arrived in New York, she felt immense pressure ...

cached_tokens is still 36 on the second call after the fix, so prompt cache reuse is unaffected; only the RNG state changed.

Testing

tests/test_sample_utils.py passes. No unit test added: the freeze only reproduces in the compiled server loop, not a direct call, so a unit test would pass on the unfixed code too. Happy to add an integration test if preferred.

Note: apply_top_k, apply_top_p, and apply_min_p carry the same random.state decorator but do not use randomness. Left unchanged here.

Fixes #1439

[nastya236]

Thanks for the fix. Can you please give tokens per second numbers with/without compile for categorical sampling?

[utkarshtiwari-24]

Thanks for the fix. Can you please give tokens per second numbers with/without compile for categorical sampling?

Thanks for your comment. Here are the numbers:

Benchmarks on an Apple M4 Pro (24GB), temp 1.0, 512 tokens, 5 runs each (generation tokens/sec):

Qwen2.5-0.5B-4bit
with compile: ~364 tokens/sec (354-376)
without (this PR): ~353 tokens/sec (350-358) -> ~3% drop

Qwen3.5-2B-4bit
with compile: ~175.8 tokens/sec (175.6-176.0)
without (this PR): ~174.8 tokens/sec (174.6-175.4) -> ~0.6% drop

Qwen3.5-9B-4bit
with compile: ~50.1 tokens/sec (50.0-50.3)
without (this PR): ~50.0 tokens/sec (49.9-50.0) -> ~0.3% drop

categorical_sampling only runs one op per token, so compiling it saves a tiny fixed amount per token. On a small model that's noticeable because each token is so fast (0.5B, ~3%), but on bigger models most of the time goes into the model itself, so it basically vanishes (~0.6% at 2B, ~0.3% at 9B). At normal model sizes it doesn't really affect speed.

[utkarshtiwari-24]

@nastya236 please review and merge the PR if approved.