adu2021 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:
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:
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:
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 tokensnum_states: Automatically determined by trie depth/breadthvocab_size: Match your tokenizer's vocabulary
Reference: Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval