Skip to content

[Engine] Metal: sampling on the decode path (temperature, top-k, top-p, min-p, penalties, Philox seed) - #4126

Open
mnott wants to merge 2 commits into
Niko1221:mainfrom
mnott:pr/metal-sampling
Open

mnott wants to merge 2 commits into
Niko1221:mainfrom
mnott:pr/metal-sampling

Conversation

@mnott

@mnott mnott commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Problem and result

The Metal decode path only ever took the argmax (metal/src/engine/window.cpp decoded greedily; sample_tokens was a stub), so temperature, top_k, top_p, min_p, the penalties and seed had no effect on macOS. In agent use the model then repeated itself word for word across turns. With this change Metal samples the way the CUDA engine does: a seed gives a reproducible text, different seeds give different texts, and greedy requests are unchanged.

Related discussion

Overview of our related macOS / Metal contributions: #4128. Focused fix of a Metal gap (the CUDA engine already samples).

Approach

  • A CPU sampler over the head's logits applies temperature, top-k, top-p, min-p and the presence/frequency/repetition penalties and draws from a Philox stream keyed by the seed (CUDA's sampler is the reference).
  • Speculative decode stays correct under sampling: a sampled verify window keeps a draft token only when the sample equals it.
  • serve: a sampled request carries a Philox seed (the request's, else a fresh random one), so temperature takes effect through the API.
  • Greedy requests (temperature 0) take the unchanged argmax path.

Validation

Branch head b5660dfcc92d5964b2527017f37c108e5da2fba1, two commits on v0.2.6. Apple Silicon, 128 GB unified memory, Gyro-M TQ2_0, --spec 2 with the MTP draft, same flags for both builds, no environment overrides.

Passed:

  • Build of strata on macOS.
  • Greedy, 80 tokens: output token ids identical to v0.2.6.
  • Sampled (--temperature 1.5 --top-k 50): seed 1 and seed 2 differ; seed 1 twice is identical.
  • Greedy decode speed, 256 tokens, 3 alternating runs: v0.2.6 63.32 / 63.26 / 63.07 tok/s, this branch 62.91 / 63.47 / 63.16 tok/s (−0.06 %).

Failed:

  • none observed

Not run:

  • CUDA, HIP and SYCL builds (the change is in the Metal engine and the server only).

Limits

The sampler runs on the CPU over the head's logits; a GPU sampler could save the logits copy per token.

mnott added 2 commits October 25, 2026 17:24
… logits: temperature, top-k, top-p, min-p, penalties, Philox seed; a sampled verify window keeps a draft when the sample equals it)
@oceanio

oceanio commented Oct 26, 2026

Copy link
Copy Markdown

This fixes the same bug (#3006) as my #3301, which has been open since 2026-10-20. Here is a comparison in case it helps the review.

I ran this PR's sample_tokens (metal/src/engine/sampler_host.cpp at b5660dfc) through #3301's parity harness. The harness checks each pick against a serial transcription of sampler.cu, on the same 248,320-logit rows, on an M5 Max.

Matches the reference. Every pick on a row without a logit_bias equals the reference's.

logit_bias is dropped on Metal. Two things cause this:

  • Verifier::init's Metal branch returns before if (le_ < 0) le_ = g.n_layers;. So set_logit_bias sees le_ == -1, takes it for an earlier pipeline stage, and never stores the bias.
  • This PR also sets sp.logit_bias = nullptr.

In the harness, 92 of 144 rows with a bias pick a different token than the reference, and all 144 pick the token they would pick without the bias. A greedy request with a bias also keeps the window's argmax, because metal_sampled does not count the bias.

Penalties are slow. The history counts are matched to the vocabulary by a linear scan for each logit:

Case Per row
Without penalties about 90 µs
penalty_last_n 64 3.8 ms
A bias and penalty_last_n 256 16 ms

The Qwen cards' default sampling has presence_penalty 1.5, and with it the server sends penalty_last_n=64. So a request that sets no sampling takes the 3.8 ms path. With --spec 2, a window samples 3 rows, which is about 11.5 ms per window. I have not measured the effect end to end.

For comparison, #3301 on the same rows takes 9-20 µs without penalties, 9.7 µs with penalty_last_n 64, and 47 µs with a bias and penalty_last_n 256. It reads the logits in place instead of copying a row.

The server seed isn't needed. The engine already decides greedy from the temperature alone (req_sp.greedy = req_temperature <= 0.0f). It also already seeds a sampled request without a seed from the clock (src/program/generate.cpp, req_sp.seed = req_seed ? req_seed : steady_clock...). So unseeded sampled requests already differ, and the serve/server.py change, which every backend runs, isn't needed for that. The bool check looks right on its own.

Two smaller points:

  • The host function now carries the CUDA kernel's name (strata::kernels::sample_tokens) and takes host pointers. With the stub gone, a later Metal caller that passes a device buffer would no longer stop at the stub.
  • std::exit(1) when there are penalties and no history ends the server. An error return through err would keep it serving.

I'm happy to fold anything useful from here into #3301, or the other way round, whichever the maintainer prefers. Since today #3301 changes no code outside #if defined(STRATA_GPU_METAL) blocks and Metal files, it is ready for review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants