# How Grammar-Constrained Decoding Works in Needle: A Technical Deep Dive

> Explore grammar constrained decoding in Needle. Learn how Needle modifies attention masks and uses KV windows and sink masks for efficient, per-token grammar rule application. Deep dive into the technical implementation.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-09-05

---

**Grammar-constrained decoding in Needle restricts generation by modifying the attention mask rather than filtering logits post-hoc, using a KV-window for sliding-window constraints and a sink mask for per-token grammar rules.**

Needle implements grammar-constrained decoding through two complementary mechanisms woven into its core attention routine. Rather than generating freely and blocking invalid tokens afterward, the decoder prevents the model from "seeing" disallowed contexts in the first place. This approach preserves efficiency and maintains full JIT compilation compatibility.

## The Two Mechanisms of Grammar-Constrained Decoding

Needle's grammar support rests on two configurable components that modify how attention operates during generation.

### KV-Window (`kv_window`): Sliding-Window Constraints

The **KV-window** limits how far back each query can attend, creating a fixed-size sliding window of allowed keys and values. For many grammars, this is sufficient—valid token sequences can be expressed as a maximum look-back length.

- **Parameter**: `kv_window` in `decode_cfg` (see [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) line 26)
- **Effect**: Tokens beyond the window receive zero attention weight
- **Use case**: Grammars where only recent context matters (e.g., n-gram constraints)

### Sink Mask (`sink`): Per-Token Grammar Rules

The **sink mask** is a Boolean tensor OR-ed into the causal mask, enabling fine-grained control over which positions each token may attend to. It supports arbitrary grammar rules like "only a verb may follow a determiner."

- **Construction**: Built in `forward_cached` (lines 68-74)
- **Application**: Injected in `_attn_cached` and `_block_cached` (lines 84-87)
- **Effect**: Disallowed tokens receive effectively -∞ attention weights

## How Grammar Constraints Flow Through Generation

The constrained decoding pipeline in `generate_cached` and `batched_generate` follows six stages:

1. **Pre-compute RoPE frequencies** via `precompute_rope_freqs`
2. **Create `DecodeCfg`** with the specified `kv_window` (line 63)
3. **Initialize KV caches** via `init_kv_cache`
4. **Build the sink mask** — default uses `_doc_prefix_len` (returns 0), but users override for grammar rules
5. **Run forward pass** (`_forward_cached`): inside `_attn_cached`, the causal mask becomes `causal = causal | sink[...]`
6. **Sample** from masked logits (`_sample`): disallowed tokens have zero probability after softmax

Because the mask is part of attention computation, invalid contexts never influence the output distribution.

## Practical Implementation Examples

### Single-Prompt Generation with KV-Window Constraints

Use `kv_window` when your grammar can be expressed as a maximum look-back distance:

```python

# Assume model_cfg, model_params, tokenizer are loaded

prompt = "The quick brown fox"

# Grammar: only allow attention to last 5 tokens

kv_window = 5

generated = needle.model.decode.generate_cached(
    config=model_cfg,
    params=model_params,
    tokenizer=tokenizer,
    prompt=prompt,
    max_new_tokens=50,
    kv_window=kv_window,
)
print(generated)

```

### Batch Generation with Custom Sink Mask

For complex grammars requiring per-token rules, construct a custom sink mask:

```python
import numpy as np

def build_grammar_sink(mask_len, prefix_len):
    """
    Example grammar: after prefix, only positions 3-5 are valid contexts
    """
    sink = np.zeros((1, mask_len), dtype=bool)
    sink[0, :prefix_len] = True           # allow prefix tokens

    sink[0, prefix_len:prefix_len+3] = True  # grammar-defined window

    return sink

sink = build_grammar_sink(mask_len=128, prefix_len=2)

texts, token_ids, log_probs = needle.model.decode.batched_generate(
    config=model_cfg,
    params=model_params,
    tokenizer=tokenizer,
    prompts=[prompt],
    max_new_tokens=30,
    kv_window=0,           # disable sliding window

    sink=sink,             # inject grammar mask

)

```

## Key Source Files

| File | Role in Grammar-Constrained Decoding |
|------|--------------------------------------|
| [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) | Core implementation: `generate_cached`, `batched_generate`, `DecodeCfg`, KV-window, sink mask injection in `_attn_cached` |
| [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) | Token-ID ↔ string conversion before/after constrained generation |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Higher-level inference entry points that forward constraint parameters |

## Performance and Architectural Benefits

Needle's approach to **grammar-constrained decoding** differs from post-hoc logit filtering in three critical ways:

- **Computational efficiency**: No separate grammar parser runs between generation steps
- **Numerical stability**: Invalid tokens receive true zero probability, not suppressed logits
- **JIT compatibility**: The entire constrained pipeline compiles with JAX/XLA

The attention-based restriction ensures that as implemented in cactus-compute/needle, grammar constraints become part of the model's fundamental computation graph.

## Summary

- **Grammar-constrained decoding** modifies attention masks, not output logits
- **KV-window** (`kv_window`) provides sliding-window constraints for look-back grammars
- **Sink mask** (`sink`) enables arbitrary per-token grammar rules via Boolean masking
- Both mechanisms operate in `_attn_cached` where `causal = causal | sink[...]` applies constraints
- The pipeline preserves full JIT compilation and avoids post-generation filtering overhead

## Frequently Asked Questions

### How is Needle's grammar-constrained decoding different from standard constrained decoding?

Standard approaches filter logits after the forward pass using a grammar parser. Needle injects constraints directly into the attention computation in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py), so invalid tokens never receive attention weight. This eliminates a separate parsing step and maintains XLA compilation.

### Can I use both KV-window and sink mask together?

Yes. Set `kv_window` to a non-zero value and pass a custom `sink` to `generate_cached` or `batched_generate`. The masks combine with OR: `causal = causal | sink[...]`. The KV-window provides coarse bounds; the sink refines with token-specific rules.

### What grammar formalisms does Needle support?

Needle itself is grammar-agnostic—it provides the attention-restriction primitives. You implement grammar logic in your `sink` construction function (see the `build_grammar_sink` example). Context-free grammars, regular expressions, or custom finite-state constraints all map to mask patterns.

### Does grammar-constrained decoding slow down generation?

No meaningful slowdown occurs because constraints live in the compiled attention kernel. According to the Needle source code, the sink mask adds only a boolean OR operation per attention head in `_attn_cached`, which XLA fuses into the surrounding computation.