Studio: apply presence_penalty on the safetensors and MLX inference paths - #6923
Conversation
…aths The safetensors and MLX generate paths resolved the inference config and then dropped presence_penalty before generation, so the same model applied the configured value under GGUF and 0 under safetensors/MLX. Thread the already-resolved presence_penalty through the orchestrator command, worker gen_kwargs, and the safetensors/MLX generate calls, and apply it with a small logits processor (subtract once per distinct completion token, prompt excluded, presence not frequency, zero is a no-op, negatives raise). Backwards compatible: presence_penalty defaults to 0.0 (byte-identical output when unset) and the GGUF path is unchanged. Also forward min_p on the legacy /generate/stream route and add the missing min_p field to GenerateRequest.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces presence penalty support for the safetensors and MLX inference paths to ensure parity with GGUF. It adds a new presence_penalty.py module with logits processors for PyTorch, implements a corresponding processor for MLX, propagates the parameter across the orchestrator, worker, and routes, and adds comprehensive unit tests. The review feedback highlights two important robustness improvements: defensively filtering out negative token IDs in PyTorch to prevent incorrect negative indexing, and safely clamping/masking indices in MLX to avoid undefined behavior or crashes from out-of-bounds tokens.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| seen = torch.unique(generated) | ||
| seen = seen[seen < vocab_size] |
There was a problem hiding this comment.
If generated contains negative token IDs (which can sometimes occur for padding or special tokens like -1 or -100), seen will also contain negative values. In PyTorch, indexing with negative values (e.g., scores[b, seen]) will index from the end of the tensor, incorrectly penalizing the last tokens in the vocabulary.
We should defensively filter out negative token IDs by ensuring seen >= 0.
| seen = torch.unique(generated) | |
| seen = seen[seen < vocab_size] | |
| seen = torch.unique(generated) | |
| seen = seen[(seen >= 0) & (seen < vocab_size)] |
There was a problem hiding this comment.
Good catch. Fixed in 87ad7b4 with exactly this: seen = seen[(seen >= 0) & (seen < vocab_size)], so a stray negative or out-of-range id is dropped before scores[b, seen] instead of wrapping to the wrong row. Real completion ids are always in range, so zero regression on the normal path.
| # Scatter-assign is idempotent for duplicate ids: presence applies once per token, on-device. | ||
| logits[:, generated] = logits[:, generated] - penalty |
There was a problem hiding this comment.
In MLX, indexing out of bounds is explicitly documented as undefined behavior. If generated contains any out-of-bounds token IDs (either negative or >= vocab_size), indexing logits[:, generated] can cause a crash or memory corruption on Apple Silicon.
Since MLX does not support boolean mask indexing for reading, we can use mx.clip to safely clamp the indices to [0, vocab_size - 1], and then multiply the penalty by a boolean mask cast to the logits' dtype to ensure out-of-bounds tokens are not penalized.
| # Scatter-assign is idempotent for duplicate ids: presence applies once per token, on-device. | |
| logits[:, generated] = logits[:, generated] - penalty | |
| import mlx.core as mx | |
| vocab_size = logits.shape[-1] | |
| in_bounds = (generated >= 0) & (generated < vocab_size) | |
| safe_generated = mx.clip(generated, 0, vocab_size - 1) | |
| # Scatter-assign is idempotent for duplicate ids: presence applies once per token, on-device. | |
| logits[:, safe_generated] = logits[:, safe_generated] - penalty * in_bounds.astype(logits.dtype) |
There was a problem hiding this comment.
Fixed in 87ad7b4. MLX disallows data-dependent shapes (boolean-mask filtering to a smaller array), so instead of clip+mask I route every out-of-range or negative id to a scratch slot at index vocab via mx.where(valid, generated, vocab), scatter-subtract into a (vocab+1)-wide penalty buffer, then apply buffer[:vocab] to the logits. Only in-range ids are penalized and nothing indexes out of bounds; the shape stays static.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a0f0d7b8d2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| return scores | ||
| vocab_size = scores.shape[-1] | ||
| for b in range(input_ids.shape[0]): | ||
| generated = input_ids[b, prompt_len:] |
There was a problem hiding this comment.
When presence_penalty > 0 and the prompt/history already contains a token, slicing from prompt_len means the first generated token is not penalized for repeating prompt tokens, and later steps only consider completion tokens. That leaves the safetensors path diverging from the GGUF/OpenAI behavior for prompts that try to discourage terms already present in context; the penalty should be based on the full token history available before sampling, not only tokens emitted after this request starts generating.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Thanks for the careful read. We checked this against the reference this PR targets. llama.cpp applies the presence penalty over a penalty_last_n window (default 64) whose ring buffer is populated by common_sampler_accept on sampled tokens only; llama-server does not seed prompt tokens into it, and our GGUF path passes presence_penalty without setting penalty_last_n. So the reference effectively penalizes generated tokens, not the prompt, which is exactly what the safetensors and MLX paths do via prompt_len. OpenAI's "tokens that appear in the text so far" wording is not defined as prompt+completion, and the stated intent is steering the response toward new topics. Basing the penalty on full history would match neither reference and would penalize every distinct token in a long user prompt, distorting logits in the long-context case. Keeping prompt exclusion, which is documented and consistent across the safetensors, MLX, and GGUF paths. Declining this one.
…both paths The presence-penalty logits processors index by generated token ids. The torch path filtered only the upper bound (seen < vocab_size), so a negative id would silently wrap to the wrong row; the MLX path had no bound at all, and MLX out-of-bounds indexing is documented undefined behavior (crash or memory corruption on Apple Silicon), unlike torch's harmless negative wrap. Bound generated ids to [0, vocab) consistently on both paths: - torch: seen[(seen >= 0) & (seen < vocab_size)] (zero-regression safety net; real completion tokens are always in range). - MLX: route out-of-range/negative ids to a discarded scratch slot via mx.where and a (vocab + 1)-wide scatter-assign mask, then subtract. MLX has no boolean-mask filtering (data-dependent output shape), so this keeps a fixed shape, stays on-device, and preserves once-per-distinct-token semantics without any torch/numpy dependency. Add torch tests for out-of-range and negative ids (only in-range distinct ids penalized, stray ids ignored, no wrong-index wrap) and a bound-documenting MLX test that runs on the arm64 macOS CI.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. 🎉 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…presence-penalty-parity
1 participant
Summary
Same-repo re-creation of #6782 (which was opened from a fork and could not merge cross-repo). Recreated onto current
main, with #6782's net changes reapplied and conflicts resolved against the refactored inference paths.The GGUF backend applies
presence_penalty(Qwen3.5/3.6 recommend1.5ininference_defaults.json), but the transformers safetensors path and the MLX path resolved the same inference config and then dropped the value before generation. The parameter was missing from the orchestrator command, the workergen_kwargs, and the safetensors/MLX generate paths, andtransformers.generate()has no nativepresence_penalty. So the same model applied the configured value as GGUF and0as safetensors/MLX, which shows up as more repetition on prompts like "Create a Flappy Bird game in HTML".This threads the already-resolved
presence_penaltythrough to the safetensors and MLX generate calls and applies it with a small logits processor. It is backwards compatible:presence_penaltydefaults to0.0(byte-identical output when unset) and the GGUF path is unchanged.Changes
core/inference/presence_penalty.py(new):apply_presence_penaltyand aLogitsProcessorfactory (OpenAI/llama.cpp semantics: subtract the penalty once from each distinct completion token, prompt excluded, presence not frequency, zero is a no-op, negatives raise). Dependency-light leaf module so the logic is unit-testable without the full backend.core/inference/inference.py: threadpresence_penaltythroughgenerate_chat_response,generate_chat_completion_with_tools, the vision path, andgenerate_stream, and attach the processor at bothmodel.generate()sites (text and vision).core/inference/orchestrator.pyandcore/inference/worker.py: carrypresence_penaltyin the generate command and workergen_kwargsso the resolved request value reaches the subprocess backend.core/inference/mlx_inference.py: matching presence-penalty logits processor for the MLX text and vision paths (on-device scatter-assign, idempotent for duplicate ids).routes/inference.py: forwardpresence_penaltyon the safetensors chat branches (tool, plain), and addpresence_penaltyplusmin_pto the legacy/generate/streamroute.models/inference.py: add the missingmin_pfield toGenerateRequest(the chat route already passedmin_p; the stream route dropped it).Tests
studio/backend/tests/test_presence_penalty.py:apply_presence_penaltyand the processor: seen token gets exactly-penalty, unseen unchanged, multiplicity ignored, negatives raise, batch rows isolated, prompt excluded, zero is a no-op, dtype/device preserved, composes with an existing processor.temperature,top_p,top_k,min_p,repetition_penalty,presence_penaltyis dropped across the route -> orchestrator command -> workergen_kwargsboundary.Supersedes #6782.