From 23c4fbfd28777df205356be4b23e988602e5a839 Mon Sep 17 00:00:00 2001 From: Oliver Nicholls Date: Thu, 10 Sep 2026 18:43:23 +0100 Subject: [PATCH 1/3] feat: pass an optional logits_processor through to llama.cpp sampling Add an optional `logits_processor` argument to `PawFunction.__call__` and thread it into the per-token `Llama.sample()` call the runtime already makes. It enables token-level constrained decoding (a regex or JSON-schema processor, for example) without reaching into `PawFunction._llm`. Defaults to `None`, which is what `Llama.sample()` already defaults to, so sampling, output parsing, and prefix caching are unchanged for existing callers. `grammar` is deliberately not exposed: `sample()` builds a temporary sampler chain per call, so a grammar passed per token would reset to its root rule on every token. --- CHANGELOG.md | 7 ++ README.md | 9 +++ docs/api-reference/python-sdk.md | 3 +- programasweights/runtime_llamacpp.py | 9 +++ tests/test_base_interpreter.py | 104 ++++++++++++++++++++++++++- 5 files changed, 129 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0190c5d..0de7d40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## Unreleased + +- Add an optional `logits_processor` argument to a compiled or base program + call. It is passed straight to llama.cpp's sampler and applied to every + generated token, enabling token-level constrained decoding (regex, + JSON schema). Defaults to `None`; sampling is unchanged when it is unset. + ## 0.4.5 (2026-09-10) - Expose structured compile API failures as `paw.APIError`, compatible with diff --git a/README.md b/README.md index 25bb9aa..ec98bf7 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,15 @@ If you need to inspect available compiler aliases programmatically, use `paw.lis GPU acceleration is enabled by default (Metal on Mac, CUDA on Linux, falls back to CPU). Set `PAW_GPU_LAYERS=0` to force CPU if GPU causes issues. +## Constrained Decoding + +A call accepts an optional `logits_processor`, passed straight to llama.cpp and applied to every generated token. Use it when the output must match a fixed shape (a regex, a JSON schema) rather than only what the spec asks for. It defaults to `None`, which samples exactly as before. + +```python +import llama_cpp +fn("Office line: +1-555-666-7777", logits_processor=llama_cpp.LogitsProcessorList([my_processor])) +``` + ## Desktop and Offline Workflows Prepare and inspect validated local assets without keeping a model loaded: diff --git a/docs/api-reference/python-sdk.md b/docs/api-reference/python-sdk.md index 01fd132..1a0af27 100644 --- a/docs/api-reference/python-sdk.md +++ b/docs/api-reference/python-sdk.md @@ -44,7 +44,7 @@ bundle directly. Required runtime metadata and base models are cached for reuse. The returned callable: ```python -output: str = fn(input_text, max_tokens=None, temperature=0.0) +output: str = fn(input_text, max_tokens=None, temperature=0.0, logits_processor=None) ``` | Parameter | Description | @@ -52,6 +52,7 @@ output: str = fn(input_text, max_tokens=None, temperature=0.0) | `input_text` | Input string for the program. | | `max_tokens` | Maximum tokens to generate. `None` (default) = use all remaining context window. | | `temperature` | Sampling temperature (default `0.0`). | +| `logits_processor` | Optional `llama_cpp.LogitsProcessorList` for token-level constrained decoding. Passed straight to llama.cpp and applied to every generated token. `None` (default) samples unconstrained. | **Context limits:** Spec + input + output share a ~2048 token window. Inputs that exceed it will error. `max_tokens` defaults to `None`: generation runs until EOS or the context limit. diff --git a/programasweights/runtime_llamacpp.py b/programasweights/runtime_llamacpp.py index 93b7ab5..5ce1cbd 100644 --- a/programasweights/runtime_llamacpp.py +++ b/programasweights/runtime_llamacpp.py @@ -397,6 +397,7 @@ def __call__( input_text: str, max_tokens: int | None = None, temperature: float = 0.0, + logits_processor: llama_cpp.LogitsProcessorList | None = None, ) -> str: """Run the program on an input. @@ -404,6 +405,10 @@ def __call__( input_text: The input to process. max_tokens: Maximum output tokens. None = use all remaining context. temperature: Sampling temperature (0 = greedy). + logits_processor: Optional llama.cpp logits processors for + token-level constrained decoding (for example a regex or + JSON-schema processor). Passed straight to llama.cpp and + applied to every generated token. None = unconstrained. Returns: The program's output as a string. @@ -444,6 +449,7 @@ def __call__( prompt_tokens, max_tokens=max_tokens, temperature=temperature, + logits_processor=logits_processor, token_description="prompt", ) @@ -459,6 +465,7 @@ def __call__( input_tokens, max_tokens=max_tokens, temperature=temperature, + logits_processor=logits_processor, prior_tokens=self._n_prefix, token_description="input", ) @@ -469,6 +476,7 @@ def _generate( *, max_tokens: int | None, temperature: float, + logits_processor: llama_cpp.LogitsProcessorList | None = None, prior_tokens: int = 0, token_description: str, ) -> str: @@ -491,6 +499,7 @@ def _generate( for _ in range(gen_limit): token = self._llm.sample( temp=temperature if temperature > 0 else 0, + logits_processor=logits_processor, ) if token == self._llm.token_eos(): diff --git a/tests/test_base_interpreter.py b/tests/test_base_interpreter.py index f56dd22..b4b1e42 100644 --- a/tests/test_base_interpreter.py +++ b/tests/test_base_interpreter.py @@ -74,6 +74,7 @@ def __init__(self, **kwargs): self.input_ids = [0] * 4096 self.tokenize_calls = [] self.eval_calls = [] + self.sample_calls = [] self.reset_calls = 0 self.closed = False self._sample_index = 0 @@ -94,12 +95,25 @@ def tokenize(self, data, *, add_bos, special): def eval(self, tokens): copied = list(tokens) self.eval_calls.append(copied) + self.input_ids[self.n_tokens:self.n_tokens + len(copied)] = copied self.n_tokens += len(copied) - def sample(self, *, temp): + def sample(self, *, temp, logits_processor=None): + self.sample_calls.append( + {"temp": temp, "logits_processor": logits_processor} + ) token = ord("A") if self._sample_index == 0 else 0 self._sample_index += 1 - return token + if logits_processor is None: + return token + # Mirror llama.cpp: every processor sees the full token + # sequence so far and returns the scores to sample from. + scores = [0.0] * 256 + scores[token] = 1.0 + input_ids = list(self.input_ids[:self.n_tokens]) + for processor in logits_processor: + scores = processor(input_ids, scores) + return max(range(len(scores)), key=scores.__getitem__) def token_eos(self): return 0 @@ -111,6 +125,7 @@ def reset(self): self.reset_calls += 1 self.n_tokens = 0 self._sample_index = 0 + self.sample_calls = [] def close(self): self.closed = True @@ -144,7 +159,11 @@ def save_prefix(_ctx, path, _seq_id, _tokens, _count): decoded_path.write_bytes(state.prefix_save_bytes) return state.prefix_save_result + class FakeLogitsProcessorList(list): + pass + module.Llama = FakeLlama + module.LogitsProcessorList = FakeLogitsProcessorList module.llama_adapter_lora_init = adapter_init module.llama_set_adapter_lora = adapter_apply module.llama_adapter_lora_free = adapter_free @@ -641,3 +660,84 @@ def worker(): assert max_active == 1 assert operations == ["open", "open"] + + +def _force_token(token: int): + """A llama.cpp logits processor that always selects one token.""" + seen = [] + + def processor(input_ids, scores): + seen.append(list(input_ids)) + forced = list(scores) + forced[token] = max(scores) + 1.0 + return forced + + processor.seen = seen + return processor + + +def test_generation_passes_no_logits_processor_by_default( + fake_runtime, + tmp_path, +): + runtime, state = fake_runtime + program_dir = _write_compiled_program(tmp_path) + + fn = runtime.PawFunction(program_dir, offline=True) + instance = state.instances[-1] + + assert fn("payload", max_tokens=1) == "A" + assert [call["logits_processor"] for call in instance.sample_calls] == [ + None + ] + fn.close() + + +def test_compiled_logits_processor_runs_on_every_generated_token( + fake_runtime, + tmp_path, +): + import llama_cpp + + runtime, state = fake_runtime + program_dir = _write_compiled_program(tmp_path) + + fn = runtime.PawFunction(program_dir, offline=True) + instance = state.instances[-1] + processor = _force_token(ord("B")) + processors = llama_cpp.LogitsProcessorList([processor]) + + output = fn("payload", max_tokens=3, logits_processor=processors) + + assert output == "BBB" + assert [call["logits_processor"] for call in instance.sample_calls] == [ + processors + ] * 3 + # One call per generated token, each seeing the full sequence so far. + assert len(processor.seen) == 3 + lengths = [len(input_ids) for input_ids in processor.seen] + assert lengths == [lengths[0], lengths[0] + 1, lengths[0] + 2] + assert [input_ids[-1] for input_ids in processor.seen[1:]] == [ + ord("B"), + ord("B"), + ] + fn.close() + + +def test_base_logits_processor_is_passed_through(fake_runtime): + import llama_cpp + + _, state = fake_runtime + _write_base_model("gpt2") + + fn = paw.function(None, interpreter="gpt2") + instance = state.instances[-1] + processor = _force_token(ord("C")) + processors = llama_cpp.LogitsProcessorList([processor]) + + assert fn("hello", max_tokens=2, logits_processor=processors) == "CC" + assert [call["logits_processor"] for call in instance.sample_calls] == [ + processors + ] * 2 + assert len(processor.seen) == 2 + fn.close() From 05228a5d9628e41975e798fdefbe205376974b7c Mon Sep 17 00:00:00 2001 From: Oliver Nicholls Date: Fri, 11 Sep 2026 16:41:22 +0100 Subject: [PATCH 2/3] fix: only forward logits_processor to sample() when it is set Passing logits_processor=None unconditionally broke callers whose sample() takes a narrower signature, including the test doubles in tests/test_local_program.py. Build the kwarg conditionally so an unset processor leaves the sampling call byte-identical to previous releases. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019XfkWkQS5DsDZmx69WBDhH --- programasweights/runtime_llamacpp.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/programasweights/runtime_llamacpp.py b/programasweights/runtime_llamacpp.py index 5ce1cbd..c2aeba6 100644 --- a/programasweights/runtime_llamacpp.py +++ b/programasweights/runtime_llamacpp.py @@ -495,11 +495,18 @@ def _generate( if prompt_tokens: self._llm.eval(prompt_tokens) + # Only forward logits_processor when one was supplied, so the call is + # byte-identical to previous releases when it is unset. Some llama.cpp + # wrappers and test doubles expose a narrower sample() signature. + sample_kwargs = {} + if logits_processor is not None: + sample_kwargs["logits_processor"] = logits_processor + output_tokens = [] for _ in range(gen_limit): token = self._llm.sample( temp=temperature if temperature > 0 else 0, - logits_processor=logits_processor, + **sample_kwargs, ) if token == self._llm.token_eos(): From ac319039a3b5e6e93d52fea6282b359776f19612 Mon Sep 17 00:00:00 2001 From: da03 Date: Sun, 13 Sep 2026 14:16:07 -0400 Subject: [PATCH 3/3] Propagate logits processor failures across native callbacks Capture processor and score-conversion errors before ctypes can swallow them, then re-raise before using the sampled token. Add native-callback regression coverage and clarify the low-level inference hook. --- CHANGELOG.md | 7 +- README.md | 5 +- docs/api-reference/python-sdk.md | 4 +- programasweights/runtime_llamacpp.py | 36 +++++-- tests/test_base_interpreter.py | 142 +++++++++++++++++++++++---- 5 files changed, 162 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0de7d40..a8b382f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,9 +3,10 @@ ## Unreleased - Add an optional `logits_processor` argument to a compiled or base program - call. It is passed straight to llama.cpp's sampler and applied to every - generated token, enabling token-level constrained decoding (regex, - JSON schema). Defaults to `None`; sampling is unchanged when it is unset. + call for caller-supplied token constraints in llama.cpp's sampler. + Defaults to `None`; sampling is unchanged when it is unset. +- Propagate processor failures to the caller instead of letting native + callback errors silently continue with unconstrained output. ## 0.4.5 (2026-09-10) diff --git a/README.md b/README.md index ec98bf7..0f05c21 100644 --- a/README.md +++ b/README.md @@ -55,13 +55,16 @@ GPU acceleration is enabled by default (Metal on Mac, CUDA on Linux, falls back ## Constrained Decoding -A call accepts an optional `logits_processor`, passed straight to llama.cpp and applied to every generated token. Use it when the output must match a fixed shape (a regex, a JSON schema) rather than only what the spec asks for. It defaults to `None`, which samples exactly as before. +A call accepts an optional `logits_processor`: an advanced hook for caller-supplied llama.cpp-compatible token constraints, not built-in regex or JSON-schema validation. It runs at every generation step. The default, `None`, keeps sampling unchanged. ```python import llama_cpp +# my_processor is your compatible callable: (input_ids, scores) -> scores. fn("Office line: +1-555-666-7777", logits_processor=llama_cpp.LogitsProcessorList([my_processor])) ``` +Processors see the full prompt and generated-token history, not just the output. Create or reset stateful processors for each call. Token limits and output whitespace trimming still apply, so validate the returned result. + ## Desktop and Offline Workflows Prepare and inspect validated local assets without keeping a model loaded: diff --git a/docs/api-reference/python-sdk.md b/docs/api-reference/python-sdk.md index 1a0af27..1b48ca7 100644 --- a/docs/api-reference/python-sdk.md +++ b/docs/api-reference/python-sdk.md @@ -52,7 +52,9 @@ output: str = fn(input_text, max_tokens=None, temperature=0.0, logits_processor= | `input_text` | Input string for the program. | | `max_tokens` | Maximum tokens to generate. `None` (default) = use all remaining context window. | | `temperature` | Sampling temperature (default `0.0`). | -| `logits_processor` | Optional `llama_cpp.LogitsProcessorList` for token-level constrained decoding. Passed straight to llama.cpp and applied to every generated token. `None` (default) samples unconstrained. | +| `logits_processor` | Optional `llama_cpp.LogitsProcessorList` of caller-supplied processors, applied at every generation step. `None` (default) keeps sampling unchanged. | + +This advanced hook is not built-in regex or JSON-schema validation. Each processor takes `(input_ids, scores)` and returns modified scores. Its token history includes the full prompt (including any compiled prefix and suffix or base-model template) plus generated tokens. Create or reset stateful processors for each call; processor exceptions propagate to the caller. Token limits and the usual output whitespace trimming still apply, so validate the returned result. **Context limits:** Spec + input + output share a ~2048 token window. Inputs that exceed it will error. `max_tokens` defaults to `None`: generation runs until EOS or the context limit. diff --git a/programasweights/runtime_llamacpp.py b/programasweights/runtime_llamacpp.py index c2aeba6..fec6d32 100644 --- a/programasweights/runtime_llamacpp.py +++ b/programasweights/runtime_llamacpp.py @@ -406,9 +406,10 @@ def __call__( max_tokens: Maximum output tokens. None = use all remaining context. temperature: Sampling temperature (0 = greedy). logits_processor: Optional llama.cpp logits processors for - token-level constrained decoding (for example a regex or - JSON-schema processor). Passed straight to llama.cpp and - applied to every generated token. None = unconstrained. + caller-supplied token constraints, applied at each generation + step with the full prompt and generated-token history. + Processor errors propagate to the caller. None = unchanged + sampling; callers must validate any structured output. Returns: The program's output as a string. @@ -495,12 +496,31 @@ def _generate( if prompt_tokens: self._llm.eval(prompt_tokens) - # Only forward logits_processor when one was supplied, so the call is - # byte-identical to previous releases when it is unset. Some llama.cpp - # wrappers and test doubles expose a narrower sample() signature. + # Leave sample() unchanged when no processors were supplied. Native + # llama.cpp invokes processors through ctypes, which swallows callback + # exceptions, so defer those errors until control returns to Python. sample_kwargs = {} + processor_error: BaseException | None = None if logits_processor is not None: - sample_kwargs["logits_processor"] = logits_processor + def guarded_processor(input_ids, scores): + nonlocal processor_error + score_count = len(scores) + try: + for processor in logits_processor: + # Match native in-place conversion, including errors + # for invalid processor return shapes or values. + scores[:] = processor(input_ids, scores) + return scores + except BaseException as exc: + processor_error = exc + # The processor may have corrupted or made scores + # read-only. Return fresh, finite logits so native sampling + # can finish; its token will be discarded below. + return [0.0] * score_count + + sample_kwargs["logits_processor"] = llama_cpp.LogitsProcessorList( + [guarded_processor] + ) output_tokens = [] for _ in range(gen_limit): @@ -508,6 +528,8 @@ def _generate( temp=temperature if temperature > 0 else 0, **sample_kwargs, ) + if processor_error is not None: + raise processor_error if token == self._llm.token_eos(): break diff --git a/tests/test_base_interpreter.py b/tests/test_base_interpreter.py index b4b1e42..baf4739 100644 --- a/tests/test_base_interpreter.py +++ b/tests/test_base_interpreter.py @@ -4,6 +4,7 @@ import hashlib import importlib import json +import math import sys import threading import types @@ -679,17 +680,23 @@ def processor(input_ids, scores): def test_generation_passes_no_logits_processor_by_default( fake_runtime, tmp_path, + monkeypatch, ): runtime, state = fake_runtime program_dir = _write_compiled_program(tmp_path) fn = runtime.PawFunction(program_dir, offline=True) instance = state.instances[-1] + temperatures = [] + + def sample_without_processor_keyword(*, temp): + temperatures.append(temp) + return ord("A") + + monkeypatch.setattr(instance, "sample", sample_without_processor_keyword) assert fn("payload", max_tokens=1) == "A" - assert [call["logits_processor"] for call in instance.sample_calls] == [ - None - ] + assert temperatures == [0] fn.close() @@ -710,34 +717,129 @@ def test_compiled_logits_processor_runs_on_every_generated_token( output = fn("payload", max_tokens=3, logits_processor=processors) assert output == "BBB" - assert [call["logits_processor"] for call in instance.sample_calls] == [ - processors - ] * 3 - # One call per generated token, each seeing the full sequence so far. - assert len(processor.seen) == 3 - lengths = [len(input_ids) for input_ids in processor.seen] - assert lengths == [lengths[0], lengths[0] + 1, lengths[0] + 2] - assert [input_ids[-1] for input_ids in processor.seen[1:]] == [ - ord("B"), - ord("B"), + assert len(instance.sample_calls) == 3 + assert processor.seen == [ + list(b"Prefix:payload:Suffix"), + list(b"Prefix:payload:SuffixB"), + list(b"Prefix:payload:SuffixBB"), ] fn.close() -def test_base_logits_processor_is_passed_through(fake_runtime): +@pytest.mark.parametrize("interpreter", ["gpt2", "Qwen/Qwen3-0.6B"]) +def test_base_logits_processor_is_passed_through(fake_runtime, interpreter): import llama_cpp _, state = fake_runtime - _write_base_model("gpt2") + _write_base_model(interpreter) - fn = paw.function(None, interpreter="gpt2") + fn = paw.function(None, interpreter=interpreter) instance = state.instances[-1] processor = _force_token(ord("C")) processors = llama_cpp.LogitsProcessorList([processor]) assert fn("hello", max_tokens=2, logits_processor=processors) == "CC" - assert [call["logits_processor"] for call in instance.sample_calls] == [ - processors - ] * 2 - assert len(processor.seen) == 2 + assert len(instance.sample_calls) == 2 + prompt = list(instance.tokenize_calls[0]["data"]) + assert processor.seen == [prompt, prompt + [ord("C")]] + fn.close() + + +def test_logits_processor_can_select_eos(fake_runtime, tmp_path): + import llama_cpp + + runtime, state = fake_runtime + fn = runtime.PawFunction(_write_compiled_program(tmp_path), offline=True) + instance = state.instances[-1] + processor = _force_token(instance.token_eos()) + eval_count = len(instance.eval_calls) + + assert fn( + "payload", + max_tokens=3, + logits_processor=llama_cpp.LogitsProcessorList([processor]), + ) == "" + assert processor.seen == [list(b"Prefix:payload:Suffix")] + assert instance.eval_calls[eval_count:] == [list(b"payload:Suffix")] + fn.close() + + +@pytest.mark.parametrize("mode", ["compiled", "base"]) +@pytest.mark.parametrize( + "exception_type,invalid_result", + [(ValueError, False), (KeyboardInterrupt, False), (SystemExit, False), + (TypeError, True)], +) +def test_logits_processor_errors_cross_native_callback_safely( + fake_runtime, tmp_path, monkeypatch, mode, exception_type, invalid_result, +): + import llama_cpp + + runtime, state = fake_runtime + if mode == "compiled": + fn = runtime.PawFunction(_write_compiled_program(tmp_path), offline=True) + expected_prompt = list(b"payload:Suffix") + else: + _write_base_model("gpt2") + fn = paw.function(None, interpreter="gpt2") + expected_prompt = list(b"payload") + instance = state.instances[-1] + error = exception_type("processor failed") + later_calls = [] + ignored_errors = [] + native_scores = [] + decoded_outputs = [] + monkeypatch.setattr(sys, "unraisablehook", ignored_errors.append) + + def failing_processor(input_ids, scores): + scores[:] = [float("-inf")] * len(scores) + if invalid_result: + # The native scores assignment must also be inside the guard. + return object() + raise error + + def later_processor(input_ids, scores): + later_calls.append(list(input_ids)) + return scores + + def native_sample(*, temp, logits_processor=None): + if logits_processor is None: + return ord("A") + scores = [0.0] * 256 + input_ids = list(instance.input_ids[:instance.n_tokens]) + + @ctypes.CFUNCTYPE(None) + def apply_processors(): + nonlocal scores + for processor in logits_processor: + scores = processor(input_ids, scores) + + apply_processors() + native_scores.append(list(scores)) + # The guarded callback must leave native sampling usable, even when + # the failed processor has already invalidated every candidate. + assert len(scores) == 256 and all(math.isfinite(score) for score in scores) + return instance.token_eos() + + def detokenize(tokens): + decoded_outputs.append(list(tokens)) + return bytes(tokens) + + monkeypatch.setattr(instance, "sample", native_sample) + monkeypatch.setattr(instance, "detokenize", detokenize) + eval_count = len(instance.eval_calls) + processors = llama_cpp.LogitsProcessorList([failing_processor, later_processor]) + with pytest.raises(exception_type) as caught: + fn("payload", max_tokens=3, logits_processor=processors) + + if not invalid_result: + assert caught.value is error + assert ignored_errors == [] + assert later_calls == [] + assert len(native_scores) == 1 + assert instance.eval_calls[eval_count:] == [expected_prompt] + assert decoded_outputs == [] + assert fn("payload", max_tokens=1) == "A" + assert len(native_scores) == 1 + assert later_calls == ignored_errors == [] fn.close()