Constrained Decoding Logit Mask

Background

JSON-mode / function-calling / grammar-constrained generation enforce structure by masking the LLM's logits: at each step, force-disable token IDs that would violate the structure. After masking, the usual argmax / sampling step is unchanged. This is the foundation of OpenAI's structured outputs, Outlines, llama.cpp grammars.

Problem statement

Implement apply_constraint_mask(logits, allowed_token_ids). Set every logit not in allowed_token_ids to -inf so downstream softmax assigns it probability 0.

Input

• logits - 1-D array of floats, length $V$ (vocabulary size).
• allowed_token_ids - iterable of int indices in $[0, V)$.

Output

• np.ndarray of shape (V,). Allowed entries unchanged; others $= -\infty$.

Examples

Input: logits=[1.0, 2.0, 3.0], allowed_token_ids={0, 2}
Output: [1.0, -inf, 3.0]

Constraints

• Do not mutate the input logits in place (return a new array).
• Empty allowed_token_ids makes every logit $-\infty$ - the caller must handle this.
• Token IDs out of range raise IndexError.

Notes

• Why -inf, not 0. Softmax of -inf is exactly 0; softmax of 0 still gets non-trivial probability.
• Grammar-constrained variants build the allowed set dynamically from a CFG state machine; this problem is the masking primitive.