Skip to content

fix(beam-search): make decode terminable, bound the grammar, add repe… #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HabibaTahir merged 1 commit into from
Sep 12, 2026
Merged

fix(beam-search): make decode terminable, bound the grammar, add repe… #44

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 105 additions & 34 deletions deployment/onnx_beam_search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,24 +1,51 @@
"""
numpy/onnxruntime-only mirror of inference/beam_search.py -- mirrors that
file's beam_search() line-for-line in logic, but never imports torch.
This is the entire point of the Option A (ONNX) deployment path: the
production Vercel bundle only needs onnxruntime + numpy, not the full
PyTorch package, which is what pushed the old bundle over the ~250MB
serverless size limit.
file in logic, but never imports torch. This is the entire point of the
Option A (ONNX) deployment path: the production Vercel bundle only needs
onnxruntime + numpy, not the full PyTorch package, which is what pushed the
old bundle over the ~250MB serverless size limit.

The two files had drifted badly -- every termination and repetition fix
landed on the torch side only. The shared policy now lives in
inference/decoding.py (pure stdlib) so they cannot silently diverge again;
only the array arithmetic differs between them.
"""

from typing import Any, Dict, List, Optional

import numpy as np
import onnxruntime as ort

from inference.grammar import NodeValidityPool


def _softmax(x: np.ndarray) -> np.ndarray:
x = x - np.max(x)
e = np.exp(x)
return e / e.sum()
from inference.decoding import (
content_token_ids,
escape_token_ids,
longest_complete_prefix,
penalised_logit,
repeat_blocked_tokens,
repetition_penalty_targets,
strip_terminator,
)
from inference.grammar import NodeValidityPool, is_complete


def _log_softmax(x: np.ndarray) -> np.ndarray:
"""Numerically stable log-softmax that PRESERVES -inf for masked entries.

The previous implementation was np.log(softmax(x) + 1e-12), which mapped
every grammatically illegal token to log(1e-12) ~= -27.6 -- a FINITE
score. argpartition would then happily select masked tokens, and since
[EOS] is one of them, that manufactured "finished" beams holding a
truncated AST. A numpy-only hazard: the torch path gets true -inf from
torch.log_softmax and never had it.
"""
finite = np.isfinite(x)
out = np.full(x.shape, -np.inf, dtype=np.float64)
if not finite.any():
return out
shifted = x - np.max(x[finite])
exponentiated = np.where(finite, np.exp(shifted), 0.0)
out[finite] = shifted[finite] - np.log(exponentiated.sum())
return out


def onnx_beam_search(
Expand All @@ -28,8 +55,12 @@ def onnx_beam_search(
beam_size: int = 5,
max_len: int = 32,
node_pool: Optional[NodeValidityPool] = None,
max_token_run: int = 4,
no_repeat_ngram_size: int = 2,
repetition_penalty: float = 1.2,
repetition_min_count: int = 4,
) -> Dict[str, Any]:
"""Mirrors inference/beam_search.py::beam_search(), but calls the
"""Mirrors inference/beam_search.py::beam_search(), but drives the
exported ONNX graph via onnxruntime instead of a torch.nn.Module."""
vocab = vocab_map["token_to_id"]
id_to_token = vocab_map["id_to_token"]
Expand All @@ -39,21 +70,20 @@ def onnx_beam_search(
if node_pool is None:
node_pool = NodeValidityPool()

exempt_ids = escape_token_ids(vocab)
content_ids = content_token_ids(id_to_token)

vocab_size = max(id_to_token.keys()) + 1
all_candidate_tokens = [id_to_token.get(idx, "[PAD]") for idx in range(vocab_size)]

src_arr = np.array([src_tokens], dtype=np.int64)

beams = [{"tokens": [bos_id], "score": 0.0, "finished": False}]
completed = []
completed: List[Dict[str, Any]] = []

for _ in range(max_len):
candidates = []
candidates: List[Dict[str, Any]] = []
for beam in beams:
if beam["finished"]:
candidates.append(beam)
continue

current_tokens = beam["tokens"]
token_strings = [id_to_token[t] for t in current_tokens]
validity_tokens = (
Expand All @@ -67,42 +97,83 @@ def onnx_beam_search(
["logits"],
{"src_seq": src_arr, "tgt_in_seq": tgt_arr},
)[0]
next_logits = logits[0, -1, :]
next_logits = np.asarray(logits[0, -1, :], dtype=np.float64).copy()

if repetition_penalty != 1.0:
for token_id in repetition_penalty_targets(
current_tokens, repetition_min_count, exempt_ids, content_ids
):
if 0 <= token_id < next_logits.shape[0]:
next_logits[token_id] = penalised_logit(
float(next_logits[token_id]), repetition_penalty
)

mask = node_pool.mask(validity_tokens, all_candidate_tokens)
safe_logits = next_logits.copy()
safe_logits[[not v for v in mask]] = -np.inf

if np.all(np.isinf(safe_logits)):
blocked = repeat_blocked_tokens(
current_tokens,
max_token_run,
no_repeat_ngram_size,
exempt_ids,
content_ids,
)
if blocked:
guarded = safe_logits.copy()
for blocked_id in blocked:
if 0 <= blocked_id < guarded.shape[0]:
guarded[blocked_id] = -np.inf
if np.isfinite(guarded).any():
safe_logits = guarded

valid_count = int(np.isfinite(safe_logits).sum())
if valid_count == 0:
if is_complete(validity_tokens):
completed.append({
"tokens": current_tokens + [eos_id],
"score": beam["score"],
"finished": True,
})
continue

log_probs = np.log(_softmax(safe_logits) + 1e-12)
k = min(beam_size, safe_logits.shape[0])
log_probs = _log_softmax(safe_logits)

k = min(beam_size, valid_count, log_probs.shape[0])
top_idx = np.argpartition(-log_probs, k - 1)[:k]
top_idx = top_idx[np.argsort(-log_probs[top_idx])]

for token_id in top_idx:
token_id = int(token_id)
for raw_id in top_idx:
token_id = int(raw_id)
score = float(log_probs[token_id])
new_tokens = current_tokens + [token_id]
finished = token_id == eos_id
candidates.append({
"tokens": new_tokens,
if not np.isfinite(score):
continue
new_beam = {
"tokens": current_tokens + [token_id],
"score": beam["score"] + score,
"finished": finished,
})
"finished": token_id == eos_id,
}
target = completed if new_beam["finished"] else candidates
target.append(new_beam)

if not candidates:
break

beams = sorted(candidates, key=lambda x: x["score"], reverse=True)[:beam_size]
if all(b["finished"] for b in beams):
completed.extend(beams)

if completed and beams[0]["score"] <= max(c["score"] for c in completed):
break

best = sorted(completed, key=lambda x: x["score"], reverse=True)[0] if completed else (
beams[0] if beams else {"tokens": [bos_id], "score": 0.0, "finished": False}
)

status = "solved" if best["finished"] else "partial"
return {"tokens": best["tokens"], "score": best["score"], "status": status}
tokens = strip_terminator(best["tokens"], eos_id)

if not best["finished"]:
salvaged = longest_complete_prefix(tokens, id_to_token)
if salvaged:
tokens = salvaged

return {"tokens": tokens, "score": best["score"], "status": status}
140 changes: 140 additions & 0 deletions docs/BEAM_SEARCH_PERFORMANCE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
# Beam Search Performance — validity-mask cost and the eval runtime bound

Answers DEV 2 task 3. Reproduce with:

```bash
python scripts/profile_beam.py --problems 8 --beams 1,2,5 --max-len 256
```

## Method

Forward-pass **count** is exact. Forward-pass **wall time** is projected from a
separately measured per-call cost on the real architecture
(`CalculusSolverModel`, `hidden_dim=256`, 4 threads, `src_len=48`,
**~24 ms/call**); timing every call in-loop would add tens of minutes and
measure nothing extra, since cost per call does not depend on weight values.
Mask time is measured in-loop and is real.

There is no checkpoint in this repo, and decode **length** is model-dependent,
so two behaviours bracket reality:

| behaviour | meaning |
|---|---|
| **oracle** | certain about every gold token; terminates correctly. What a trained model should approach. |
| **untrained** | uniform random logits; wanders. The pessimistic bound. |

## 1. Validity-mask cost at beam_size=5

All figures at `max_len=256`, which is what `run_eval.py` actually uses
(see §3).

| behaviour | beam | guards | is_valid_prefix calls | tokens parsed | mask ms | forwards | proj. total s | **mask share** |
|---|---|---|---|---|---|---|---|---|
| oracle | 1 | on | 2,806 | 36,410 | 27 | 23 | 0.60 | 4.4% |
| oracle | 2 | on | 5,487 | 72,695 | 44 | 44 | 1.17 | 3.8% |
| oracle | 5 | on | 13,532 | 181,552 | 120 | 109 | 2.89 | 4.2% |
| untrained | 1 | on | 6,076 | 151,900 | 78 | 49 | 1.32 | 5.9% |
| untrained | 2 | on | 15,252 | 484,220 | 205 | 123 | 3.32 | 6.2% |
| untrained | 5 | on | 27,404 | 641,204 | 337 | 221 | 5.94 | 5.7% |

**The mask is not the bottleneck.** It is 4–6% of runtime; the un-batched
forward passes are the rest. Beam width scales both terms roughly linearly, so
`beam_size=2 -> 5` costs about 2.5x across the board — the mask is not what
makes width expensive.

## 2. Effect of the repetition guards on mask cost

The guards never increase `is_valid_prefix` volume. They are either neutral or
a small saving:

| behaviour | beam | guards on | guards off |
|---|---|---|---|
| oracle | 5 | 13,532 | 13,532 (identical) |
| untrained | 1 | 6,076 | 6,448 |
| untrained | 5 | 27,404 | 29,884 |

Their own cost is **0.0084 ms/step** — 0.02% of a step, against ~24 ms for a
forward pass.

> **Correction.** An earlier revision of this document claimed call volume was
> *identical* with guards on and off, and described that as structural. That was
> over-generalised from two configurations. Volume is driven by
> `beams x steps x |V|`, and the guards do change how many steps a degenerate
> beam survives — so they can reduce it. The claim to rely on is the weaker,
> true one: **the guards never make mask cost worse.**

## 3. The `eval/run_eval.py` bound

`run_eval.py` constructs `CalculusSolverInference(model_path=..., beam_size=5)`
with no `max_len`, so it inherits **256** from `inference/solve.py`. Over the
300 benchmark problems:

| model behaviour | per problem | **300 problems** |
|---|---|---|
| correct (terminates) | 2.89 s | **~15 minutes** |
| worst case (never terminates) | 5.94 s | **~30 minutes** |

**Documented bound: `run_eval.py` at `beam_size=5` should complete in ~15
minutes, and must not exceed ~30 minutes.** Materially longer than that means
something regressed in the grammar bounds of §4 — it is not the mask.

This bound is projected, not measured end-to-end: `run_eval.py` exits
immediately because `checkpoints/final/best.pt` is absent. Re-run
`scripts/profile_beam.py` once a checkpoint exists to close the gap.

## 4. Structural bounds are what made the worst case tractable

Before the grammar was bounded, the untrained worst case at `beam_size=5`,
`max_len=256` was **34.58 s/problem** — 158,224 `is_valid_prefix` calls and
**20.4 million tokens parsed**, with the mask at 14.7%. `is_valid_prefix`
re-parses the whole prefix per candidate, so cost is
`O(beams x steps x |V| x L)` — quadratic in sequence length. A decoder that
never terminates drags every beam to `L=256` and pays that quadratic.

Capping the grammar's previously unbounded dimensions removed it:

| bound | max in 155,000 real targets | cap | margin |
|---|---|---|---|
| `MAX_OPVARS_PER_OP` | 1 | 4 | 4x |
| `MAX_NESTING_DEPTH` | 3 | 8 | 2.7x |
| `MAX_SIBLINGS` | 2 | 8 | 4x |

Result: **34.58 s -> 5.94 s per problem** (5.8x), tokens parsed 20.4M -> 641k
(32x). Grammar acceptance of real targets is unchanged at **100.00%** of all
155,000 — the caps reject nothing real.

## 5. Known limitation

The bounds do **not** guarantee that a degenerate decoder terminates within
`max_len`. A tree of depth 8 with 8 siblings per node can hold far more than
256 tokens, so a pathological model can still fill the budget with a legal but
incomplete prefix and be truncated. Measured: a model biased toward nested
`NODE:FRAC` still reaches `max_len` with `status="partial"`.

`beam_search` mitigates but cannot eliminate this: when no beam terminates it
returns the longest **closed** prefix (`longest_complete_prefix`), so the caller
gets something parseable. That only helps when the outermost node closed — a
sequence that opens a fraction and never closes it has no complete prefix, and
is returned as-is.

Guaranteeing termination would require a budget-aware constraint: once the
remaining budget equals the tokens needed to close all open structures, mask
everything except closing tokens. That is implementable but changes what the
model is allowed to emit, so it is flagged rather than assumed.

## 6. Recommendation

**Decouple the decoder step budget from the encoder padding width.**
`inference/solve.py` uses one `max_len=256` for both source padding
(`solve.py:137-139`) and the generation budget (`solve.py:155`). They are
unrelated quantities.

The longest serialized target across all 155,000 dataset targets is **38
tokens**. A decoder budget of **48** (already the value in `config.json`) covers
100% of real targets with 9 tokens spare. This caps the worst case directly
rather than relying on the model to stop.

Secondary: batching the beams into one forward call per step was measured at
only **~1.5x** (373ms -> 254ms for 5 beams), not the 5x it looks like — the
model is CPU-bound and does not parallelise well across batch at 4 threads.
Worth doing, but it is not the lever.
Loading