diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 064be30c06..167706f701 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -31,6 +31,7 @@ from core.inference.chat_eos import ( chat_eos_repair, resolve_chat_turn_end_eos_ids_using, ) +from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog from loggers import get_logger @@ -832,6 +833,7 @@ class InferenceBackend: tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + presence_penalty: float = 0.0, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -865,6 +867,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) initial = list(messages) @@ -901,12 +904,14 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response for text or vision models (lock held by background thread). ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking`` are forwarded into ``apply_chat_template`` so templates that understand them (Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls. + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_chat_response_inner( messages = messages, @@ -923,6 +928,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_chat_response_inner( @@ -942,6 +948,7 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic, called by generate_chat_response and generate_with_adapter_control. @@ -981,6 +988,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = cancel_event, + presence_penalty = presence_penalty, ) return else: @@ -1093,6 +1101,7 @@ class InferenceBackend: repetition_penalty, cancel_event = cancel_event, _adapter_state = _adapter_state, + presence_penalty = presence_penalty, ) def _generate_vision_response( @@ -1107,6 +1116,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] @@ -1196,6 +1206,14 @@ class InferenceBackend: top_k = top_k, min_p = min_p, ) + # Presence penalty (GGUF parity) for VLM chat. + _vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None + if _vision_input_ids is not None: + _pp = _make_presence_penalty_processor( + presence_penalty, int(_vision_input_ids.shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp err: dict[str, str] = {} @@ -1424,11 +1442,13 @@ class InferenceBackend: repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). _adapter_state: if not None, the background thread toggles adapters before model.generate(), under _generation_lock. + ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: yield "Error: No active model" @@ -1489,6 +1509,12 @@ class InferenceBackend: if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. + _pp = _make_presence_penalty_processor( + presence_penalty, int(inputs["input_ids"].shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp if cancel_event is not None: from transformers.generation.stopping_criteria import ( StoppingCriteria, diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 45f46fef2f..62d268e15f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -41,6 +41,50 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _make_mlx_presence_penalty_processor(penalty: float): + """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. + + generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the + full running sequence; the first call is prompt-only, so latch that length + and penalize only after it. + """ + state = {"prompt_len": None} + + def _processor(tokens, logits): + if state["prompt_len"] is None: + # First call = prompt only; latch its length. + state["prompt_len"] = int(tokens.shape[0]) + return logits + generated = tokens[state["prompt_len"] :] + if generated.size == 0: + return logits + import mlx.core as mx + + vocab = logits.shape[-1] + # Bound generated ids to the valid range [0, vocab) before they index + # logits. MLX does no bounds checking and out-of-bounds indexing is + # documented undefined behavior (crash / memory corruption), unlike the + # torch path's harmless negative wrap -- so this bound is load-bearing + # here and matches the torch filter seen[(seen >= 0) & (seen < vocab)]. + # MLX has no boolean-mask filtering (data-dependent output shape is + # unsupported), so instead of compacting the id list we route every + # out-of-range or negative id to a scratch slot at index ``vocab`` that + # is dropped before the subtract. That scratch slot can never collide + # with a real token, so real ids (including id 0) are penalized exactly + # once and stray ids are ignored. + valid = (generated >= 0) & (generated < vocab) + safe = mx.where(valid, generated, vocab).astype(mx.int32) + # Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate + # ids are idempotent, so presence applies once per distinct token; the + # scratch column is discarded and the full-width subtract stays on-device. + mask = mx.zeros((vocab + 1,), dtype = logits.dtype) + mask[safe] = penalty + logits = logits - mask[:vocab] + return logits + + return _processor + + class MLXInferenceBackend: def __init__(self): self.models = {} @@ -282,6 +326,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -329,6 +374,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) else: yield from self._generate_text( @@ -344,6 +390,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_text( @@ -361,6 +408,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors @@ -407,15 +455,21 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Only build a logits processor for a non-trivial repetition penalty. - logits_processors = None + # Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths). + logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, ): - logits_processors = make_logits_processors( - repetition_penalty = float(repetition_penalty), + logits_processors.extend( + make_logits_processors( + repetition_penalty = float(repetition_penalty), + ) ) + if presence_penalty: + logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + if not logits_processors: + logits_processors = None token_ids = [] logger.info( @@ -481,6 +535,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_vlm import stream_generate as vlm_stream @@ -528,10 +583,23 @@ class MLXInferenceBackend: top_k = int(top_k or 0), min_p = float(min_p or 0.0), ) - if repetition_penalty is not None and float(repetition_penalty) not in ( + _rep_active = repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, - ): + ) + if presence_penalty: + # Presence needs a custom processor: pass the full list (repetition + + # presence) instead of the repetition_penalty shortcut so both apply once. + from mlx_lm.sample_utils import make_logits_processors + + _vlm_processors = [] + if _rep_active: + _vlm_processors.extend( + make_logits_processors(repetition_penalty = float(repetition_penalty)) + ) + _vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + vlm_kwargs["logits_processors"] = _vlm_processors + elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) with self._generation_lock: diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 394ff97eac..19d2230278 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -428,6 +428,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> dict: """Build the 'generate' command shared by the locked and dispatched paths.""" cmd = { @@ -442,6 +443,7 @@ class InferenceOrchestrator: "min_p": min_p, "max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty, + "presence_penalty": presence_penalty, } # Only forward template kwargs the caller set, for older worker compat. if use_adapter is not None: @@ -631,6 +633,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -684,6 +687,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -1166,6 +1170,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess. @@ -1175,6 +1180,8 @@ class InferenceOrchestrator: ``stats_holder``: caller-owned dict; on gen_done its "stats" key gets the worker's usage/timings. Request-scoped to avoid cross-stream reads. + + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_inner( messages = messages, @@ -1193,6 +1200,7 @@ class InferenceOrchestrator: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, stats_holder = stats_holder, + presence_penalty = presence_penalty, ) def generate_chat_completion_with_tools( @@ -1220,6 +1228,7 @@ class InferenceOrchestrator: bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -1255,6 +1264,7 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, # last turn wins, like the GGUF tool loop stats_holder = stats_holder, + presence_penalty = presence_penalty, ) if use_adapter is not None: yield from self.generate_with_adapter_control( @@ -1322,6 +1332,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -1365,6 +1376,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, diff --git a/studio/backend/core/inference/presence_penalty.py b/studio/backend/core/inference/presence_penalty.py new file mode 100644 index 0000000000..c73c513887 --- /dev/null +++ b/studio/backend/core/inference/presence_penalty.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Presence-penalty logits helpers for the safetensors/MLX inference paths. + +Kept in a dependency-light leaf module (torch + transformers only, no unsloth / +peft) so the pure logic can be imported and unit-tested without pulling in the +full inference backend. ``core.inference.inference`` re-exports these for the +runtime generate paths. +""" + +import torch + + +def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int): + """OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct + completion token (positions >= prompt_len; prompt excluded, multiplicity + ignored, negatives raise). In place; zero is a no-op.""" + if not penalty: + return scores + vocab_size = scores.shape[-1] + for b in range(input_ids.shape[0]): + generated = input_ids[b, prompt_len:] + if generated.numel() == 0: + continue + seen = torch.unique(generated) + # Bound generated ids to the valid range [0, vocab_size). Real completion + # tokens are always in range, so this is a zero-regression safety net that + # drops any stray out-of-range or negative id before indexing (mirrors the + # MLX path's bound). Filtering both ends avoids indexing scores with a + # negative id (which would silently wrap to the wrong row). + seen = seen[(seen >= 0) & (seen < vocab_size)] + if seen.numel(): + scores[b, seen] = scores[b, seen] - penalty + return scores + + +def _make_presence_penalty_processor(penalty: float, prompt_len: int): + """``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical).""" + if not penalty: + return None + from transformers import LogitsProcessor, LogitsProcessorList + + class _PresencePenaltyLogitsProcessor(LogitsProcessor): + @torch.no_grad() + def __call__(self, input_ids, scores): + return apply_presence_penalty(input_ids, scores, penalty, prompt_len) + + return LogitsProcessorList([_PresencePenaltyLogitsProcessor()]) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 615c5c5a0d..d4b102e422 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -457,6 +457,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: "min_p": cmd.get("min_p", 0.0), "max_new_tokens": cmd.get("max_new_tokens", 256), "repetition_penalty": cmd.get("repetition_penalty", 1.0), + "presence_penalty": cmd.get("presence_penalty", 0.0), "cancel_event": cancel_event, } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 1e7770a7f8..0f27b695fe 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -177,6 +177,7 @@ class GenerateRequest(BaseModel): temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") + min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling") max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3f6f6a2a78..5332037e0d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4353,8 +4353,10 @@ async def generate_stream( temperature = request.temperature, top_p = request.top_p, top_k = request.top_k, + min_p = request.min_p, max_new_tokens = request.max_new_tokens, repetition_penalty = request.repetition_penalty, + presence_penalty = request.presence_penalty, cancel_event = cancel_event, ) _DONE = object() @@ -6995,6 +6997,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, @@ -7240,6 +7243,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, ) # Forward reasoning kwargs; the worker/template wrapper peels off any the # template doesn't accept. diff --git a/studio/backend/tests/test_presence_penalty.py b/studio/backend/tests/test_presence_penalty.py new file mode 100644 index 0000000000..030ddb6011 --- /dev/null +++ b/studio/backend/tests/test_presence_penalty.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths. + +The safetensors path historically dropped ``presence_penalty``, so the SAME model +looked worse served as safetensors. These tests pin the processor semantics +(subtract once per distinct completion token, prompt excluded, presence not +frequency, zero a no-op, negatives raise) plus a param-propagation regression +over route -> orchestrator cmd -> worker gen_kwargs. +""" + +import threading + +import pytest +import torch + +from core.inference.presence_penalty import ( + apply_presence_penalty, + _make_presence_penalty_processor, +) + + +def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged(): + input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2) + assert out[0, 3].item() == pytest.approx(-1.5) + for tok in (0, 1, 2, 4): + assert out[0, tok].item() == pytest.approx(0.0) + + +def test_multiplicity_ignored_presence_not_frequency(): + # Token 3 emitted three times -> still a single -penalty (presence, not freq). + input_ids = torch.tensor([[0, 3, 3, 3]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1) + assert out[0, 3].item() == pytest.approx(-2.0) + + +def test_negative_penalty_raises_seen_logits(): + input_ids = torch.tensor([[0, 2]]) + scores = torch.zeros(1, 4) + out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1) + assert out[0, 2].item() == pytest.approx(0.5) + + +def test_prompt_tokens_excluded(): + # Token 7 is prompt-only (untouched); token 4 in the completion is penalized. + input_ids = torch.tensor([[7, 4, 4]]) + scores = torch.zeros(1, 8) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 7].item() == pytest.approx(0.0) + assert out[0, 4].item() == pytest.approx(-1.0) + + +def test_batch_rows_isolated(): + input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2] + scores = torch.zeros(2, 4) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 1].item() == pytest.approx(-1.0) + assert out[0, 2].item() == pytest.approx(0.0) + assert out[1, 2].item() == pytest.approx(-1.0) + assert out[1, 1].item() == pytest.approx(0.0) + + +def test_zero_penalty_is_noop(): + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1) + assert torch.equal(out, original) + + +def test_empty_completion_is_noop(): + # prompt_len covers the whole sequence -> nothing generated yet. + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3) + assert torch.equal(out, original) + + +def test_out_of_vocab_id_ignored(): + # A generated id >= vocab_size (defensive) must not index out of bounds. + input_ids = torch.tensor([[0, 9]]) + scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_negative_generated_id_ignored(): + # A negative generated id (defensive) must be dropped, not wrap to scores[-1]. + input_ids = torch.tensor([[0, -1]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + # Nothing penalized; in particular the last row (the numpy/torch wrap target + # for id -1) is untouched. + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized(): + # Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a + # negative id (-1). Only the in-range distinct id is penalized; OOB/negative + # ids are ignored with no crash and no wrong-index wrap. This fails under the + # old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and + # passes only with the both-ends bound. + input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + expected = torch.zeros(1, 5) + expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored) + assert torch.equal(out, expected) + assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row + + +def test_dtype_and_device_preserved(): + input_ids = torch.tensor([[0, 1]]) + scores = torch.zeros(1, 4, dtype = torch.float16) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out.dtype == torch.float16 + assert out.device == scores.device + + +def test_processor_none_when_zero(): + assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None + + +def test_processor_applies_penalty(): + proc = _make_presence_penalty_processor(1.5, prompt_len = 2) + assert proc is not None + input_ids = torch.tensor([[0, 1, 3]]) + scores = torch.zeros(1, 5) + out = proc(input_ids, scores) + assert out[0, 3].item() == pytest.approx(-1.5) + + +def test_processor_composes_with_other_processors(): + # LogitsProcessorList must run our processor alongside a pre-existing one. + from transformers import LogitsProcessor, LogitsProcessorList + + class _AddToTokenZero(LogitsProcessor): + def __call__(self, input_ids, scores): + scores[:, 0] = scores[:, 0] + 100.0 + return scores + + presence = _make_presence_penalty_processor(1.0, prompt_len = 1) + combined = LogitsProcessorList([_AddToTokenZero(), *presence]) + input_ids = torch.tensor([[5, 2]]) # completion = [2] + scores = torch.zeros(1, 6) + out = combined(input_ids, scores) + assert out[0, 0].item() == pytest.approx(100.0) # other processor ran + assert out[0, 2].item() == pytest.approx(-1.0) # presence ran + + +def test_mlx_presence_penalty_callable(): + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.5) + # First call = prompt only (latches prompt_len, penalizes nothing). + prompt = mx.array([10, 11]) + logits0 = mx.zeros((1, 20)) + out0 = proc(prompt, logits0) + assert float(out0[0, 10]) == pytest.approx(0.0) + # Second call: one completion token (5) appended -> penalized once. + seq = mx.array([10, 11, 5]) + logits1 = mx.zeros((1, 20)) + out1 = proc(seq, logits1) + assert float(out1[0, 5]) == pytest.approx(-1.5) + assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched + + +def test_mlx_presence_penalty_bounds_out_of_range_ids(): + # Documents (and, on Apple Silicon CI, enforces) the intended MLX bound: + # out-of-vocab and negative completion ids must be ignored. MLX does no + # bounds checking and OOB indexing is undefined behavior (crash / memory + # corruption), so the processor routes stray ids to a discarded scratch slot + # and penalizes only in-range distinct ids -- matching the torch filter + # seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent. + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.0) + proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2 + # Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a + # negative id (-1); only the in-range id is penalized and nothing crashes. + seq = mx.array([10, 11, 3, 99, -1]) + out = proc(seq, mx.zeros((1, 8))) + assert float(out[0, 3]) == pytest.approx(-1.0) + for tok in range(8): + if tok != 3: + assert float(out[0, tok]) == pytest.approx(0.0) + + +# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs +_SAMPLING = { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.05, + "repetition_penalty": 1.1, + "presence_penalty": 1.5, +} + + +def test_orchestrator_cmd_carries_all_sampling_params(): + from core.inference.orchestrator import InferenceOrchestrator + + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + cmd = o._build_generate_cmd( + "req1", + None, + messages = [{"role": "user", "content": "hi"}], + max_new_tokens = 128, + **_SAMPLING, + ) + for key, val in _SAMPLING.items(): + assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd" + + +def test_worker_forwards_all_sampling_params_to_backend(): + from core.inference.worker import _handle_generate + + class _RecordingBackend: + last_generation_stats = None + + def __init__(self): + self.received = None + + def generate_chat_response(self, **kwargs): + self.received = kwargs + return iter(()) # empty stream -> loop exits, gen_done is sent + + class _FakeQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + cmd = { + "type": "generate", + "request_id": "r", + "messages": [{"role": "user", "content": "hi"}], + "max_new_tokens": 128, + **_SAMPLING, + } + backend = _RecordingBackend() + _handle_generate(backend, cmd, _FakeQueue(), threading.Event()) + + assert backend.received is not None + for key, val in _SAMPLING.items(): + assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs"