Block Unsafe Tokens Before Generation with resk-logits — GPU-Accelerated Aho-Corasick for LLM Safety

Links

GitHub: https://github.com/Resk-Security/resk-logits
PyPI: https://pypi.org/project/resklogits
Site: https://resk.fr

Most LLM safety filters are reactive: they scan generated text for bad patterns after the model already output them. For production systems, thats too late. The token is already sampled, logged, and potentially served to a user.

resk-logits takes a different approach: it intercepts at the logits level, before the model samples a token. Using a GPU-accelerated Aho-Corasick automaton, it matches 10000+ unsafe patterns against potential token completions in real time and applies a configurable penalty to matching logits.

How it works

import torch
from resklogits import LogitsProcessor

processor = LogitsProcessor()

# Load your patterns
processor.add_patterns([
    "ignore previous instructions",
    "role: system",
    "<|im_end|>",
    # ... more patterns
])

# In your generation loop
class ReskLogitsWarper:
    def __call__(self, input_ids, scores):
        return processor(scores, mode="hard")

# Use with HuggingFace generate
outputs = model.generate(
    input_ids,
    logits_processor=[ReskLogitsWarper()]
)

Key features

GPU-accelerated: runs 10000+ patterns in under 1ms on RTX 4090
Two severity modes: hard sets matching logits to -inf, bias applies a configurable penalty
Dynamic reloading: add or remove patterns at runtime
PyTorch native: works with any HuggingFace model
C++/CUDA backend: minimal overhead in the generation loop