feat: pass an optional logits_processor through to llama.cpp sampling

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## Unreleased
+
+- Add an optional `logits_processor` argument to a compiled or base program
+  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)
 
 - Expose structured compile API failures as `paw.APIError`, compatible with

README.md

@@ -53,6 +53,18 @@ 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`: 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:

docs/api-reference/python-sdk.md

@@ -44,14 +44,17 @@ 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 |
 |-----------|-------------|
 | `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` 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.
 

programasweights/runtime_llamacpp.py

@@ -397,13 +397,19 @@ 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.
 
         Args:
             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
+                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.
@@ -444,6 +450,7 @@ def __call__(
                 prompt_tokens,
                 max_tokens=max_tokens,
                 temperature=temperature,
+                logits_processor=logits_processor,
                 token_description="prompt",
             )
 
@@ -459,6 +466,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 +477,7 @@ def _generate(
         *,
         max_tokens: int | None,
         temperature: float,
+        logits_processor: llama_cpp.LogitsProcessorList | None = None,
         prior_tokens: int = 0,
         token_description: str,
     ) -> str:
@@ -487,11 +496,40 @@ def _generate(
         if prompt_tokens:
             self._llm.eval(prompt_tokens)
 
+        # 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:
+            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):
             token = self._llm.sample(
                 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

tests/test_base_interpreter.py

@@ -4,6 +4,7 @@
 import hashlib
 import importlib
 import json
+import math
 import sys
 import threading
 import types
@@ -74,6 +75,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 +96,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 +126,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 +160,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 +661,185 @@ 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,
+    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 temperatures == [0]
+    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 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()
+
+
+@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(interpreter)
+
+    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 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()