Static Constrained Decoding

---
name: static-constrained-decoding
title: "Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2602.22647"
keywords: [Constrained Decoding, Trie Vectorization, Generative Retrieval, Inference Optimization, GPU Acceleration]
description: "STATIC converts prefix trees into sparse matrices for vectorized constrained decoding, achieving 948x speedup over CPU and enabling production-scale recommendation systems."
---

# Technique: Sparse Matrix Vectorization for Constrained Decoding

Generative retrieval systems output item IDs from language models to recommend products, videos, or content. However, these systems must respect business constraints: items must be fresh, appropriate for user age, in stock, or in specific categories. Traditional prefix-tree (trie) based constrained decoding incurs massive latency penalties on GPUs because tree traversal is inherently sequential—each token generation requires tree navigation, incompatible with hardware vectorization.

STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding) solves this by flattening trie structures into Compressed Sparse Row (CSR) matrices. This converts sequential tree traversal into vectorized sparse matrix operations that GPUs can parallelize, achieving production-scale performance with negligible overhead.

## Core Concept

The core insight: prefix trees are transition structures that can be represented as sparse matrices. Rather than traversing a tree sequentially, convert it to a static sparse matrix where each row represents valid next-token constraints. During generation, each step becomes a sparse matrix multiplication instead of tree node lookup.

This enables hardware-native parallelization: GPUs excel at sparse matrix operations but struggle with tree traversal. By changing the representation, STATIC unlocks GPU efficiency for constrained decoding.

## Architecture Overview

- **Trie to CSR Conversion**: Flatten constraint tree into static sparse matrix
- **Sparse Matrix Ops**: Use GPU-native sparse-matrix-multiply for token filtering
- **Stateless Decoding**: No tree state needed during generation; matrix row index suffices
- **Minimal Overhead**: Single sparse operation per token (0.033 ms on GPU)
- **Production Deployment**: Tested at scale on video recommendation platform

## Implementation Steps

STATIC involves converting constraint tries to sparse matrices and using them during inference. Here's how to implement it:

Build a constraint trie from a set of allowed item IDs. In practice, this represents catalog items you can recommend:

```python
import numpy as np
from scipy.sparse import csr_matrix

class TrieToSparseMatrix:
    def __init__(self, vocab_size):
        self.vocab_size = vocab_size
        self.trie = {}
        self.row_idx = 0
        self.rows = []
        self.cols = []
        self.data = []

    def insert_item_id(self, item_id_str):
        """
        Add a single item ID (string) to the constraint set.
        Example: item_id_str = "product_12345" or "video_98765"
        """
        tokenized = self.tokenize(item_id_str)  # Convert string to token sequence
        self.insert_sequence(tokenized)

    def insert_sequence(self, token_sequence):
        """Insert a sequence of tokens into trie."""
        current = self.trie
        for token in token_sequence:
            if token not in current:
                current[token] = {}
            current = current[token]
        current['_end'] = True  # Mark end-of-sequence

    def tokenize(self, item_id_str):
        """Convert item ID string to token IDs."""
        # In practice, use your tokenizer
        # For now, assume byte-level tokenization
        return [ord(c) for c in item_id_str[:20]]  # Limit to 20 bytes

    def trie_to_csr_matrix(self):
        """
        Convert trie to sparse matrix where rows are states,
        columns are tokens, and values indicate valid transitions.
        """
        state_to_idx = {}
        state_idx = 0

        # BFS through trie to assign state IDs
        queue = [('root', self.trie)]
        state_to_idx['root'] = 0
        state_idx = 1

        while queue:
            state_name, state_dict = queue.pop(0)
            current_idx = state_to_idx[state_name]

            for token, next_state in state_dict.items():
                if token == '_end':
                    continue
                if isinstance(next_state, dict):
                    state_name_next = f"{state_name}_{token}"
                    if state_name_next not in state_to_idx:
                        state_to_idx[state_name_next] = state_idx
                        state_idx += 1
                        queue.append((state_name_next, next_state))

                    # Add transition: (current_state, token) -> next_state
                    next_idx = state_to_idx[state_name_next]
                    self.rows.append(current_idx)
                    self.cols.append(token)
                    self.data.append(next_idx)

        # Build CSR matrix
        csr = csr_matrix(
            (self.data, (self.rows, self.cols)),
            shape=(state_idx, self.vocab_size),
            dtype=np.int32
        )
        return csr, state_to_idx

    def tokenize(self, text):
        return [ord(c) % self.vocab_size for c in text]
```

Use the sparse matrix during token generation to enforce constraints:

```python
class ConstrainedDecoder:
    def __init__(self, model, tokenizer, csr_matrix, num_states):
        self.model = model
        self.tokenizer = tokenizer
        self.csr_matrix = csr_matrix
        self.num_states = num_states
        self.current_state = 0  # Start at root of trie

    def generate_constrained(self, prompt, max_length=50):
        """
        Generate tokens while respecting constraints from sparse matrix.
        """
        input_ids = self.tokenizer.encode(prompt)
        self.current_state = 0  # Reset to root

        for _ in range(max_length):
            # Get model logits
            with torch.no_grad():
                outputs = self.model(torch.tensor([input_ids]))
                logits = outputs.logits[0, -1, :]  # Last token's logits

            # Get valid next tokens from CSR matrix
            row_data = self.csr_matrix.getrow(self.current_state)
            valid_tokens = row_data.nonzero()[1]  # Column indices = valid tokens

            if len(valid_tokens) == 0:
                # End of constraint path reached
                break

            # Mask logits: zero out invalid tokens
            masked_logits = logits.clone()
            all_tokens = set(range(len(logits)))
            invalid_tokens = all_tokens - set(valid_tokens.tolist())
            masked_logits[list(invalid_tokens)] = -float('inf')

            # Sample from valid tokens only
            next_token = torch.argmax(masked_logits).item()
            input_ids.append(next_token)

            # Update state
            next_state_options = row_data[0, next_token]
            if next_state_options > 0:
                self.current_state = next_state_options
            else:
                break

        return self.tokenizer.decode(input_ids)
```

Integrate into recommendation pipeline:

```python
def create_recommendation_decoder(item_catalog, model, tokenizer, vocab_size):
    """
    Setup constrained decoder for product recommendation.
    """
    # Convert item IDs to constraint trie
    trie_converter = TrieToSparseMatrix(vocab_size)
    for item_id in item_catalog:
        trie_converter.insert_item_id(item_id)

    # Build sparse matrix
    csr, state_dict = trie_converter.trie_to_csr_matrix()

    # Create decoder
    decoder = ConstrainedDecoder(model, tokenizer, csr, len(state_dict))
    return decoder

# Usage
decoder = create_recommendation_decoder(
    item_catalog=['video_123', 'video_456', 'video_789'],
    model=model,
    tokenizer=tokenizer,
    vocab_size=50257
)

recommendation = decoder.generate_constrained("User watches technology videos. Recommend:", max_length=20)
```

## Practical Guidance

**When to Use:**
- Generative retrieval systems (item/document recommendation)
- Catalog constraints (freshness, category, inventory)
- Any system requiring token-level output restrictions
- Production deployments where latency is critical

**When NOT to Use:**
- Open-ended text generation (no constraints)
- Extremely large constraint sets (billions of items) where memory is limited
- Systems where constraint changes frequently (trie conversion is offline)

**Performance:**
- CPU implementation: ~1–5ms per token (slow)
- GPU implementation: ~0.033ms per token (fast)
- Speedup: 948x over CPU baselines
- Scales to production catalogs (millions of items)

**Implementation Notes:**
- Precompute CSR matrix once; reuse across requests
- Use GPU-native sparse operations (cuSPARSE on NVIDIA, HIP on AMD)
- State index fits in int32 for most practical catalogs
- Handle end-of-sequence detection to avoid incomplete items

**Hyperparameters:**
- `max_length`: Depends on item ID length; typical 10–30 tokens
- `num_states`: Automatically determined by trie depth/breadth
- `vocab_size`: Match your tokenizer's vocabulary

---

**Reference:** [Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval](https://arxiv.org/abs/2602.22647)