DEV Community

RESK
RESK

Posted on

GPU accelerated Aho-Corasick logits processor for LLM jailbreak blocking with resk-logits

#python #llm #cybersecurity #opensource

Animated Aho-Corasick automaton scanning the LLM vocabulary — watch patterns being matched in real time across ~50k tokens.

resk-logits Aho-Corasick GPU animation

Links

What is resk-logits?

resk-logits is a GPU accelerated Aho-Corasick logits processor for LLM safety. It operates before the model samples a token — on the logits distribution itself — rather than scanning generated output text after the fact. This means zero latency overhead on generation and guaranteed blocking even when the model is under adversarial pressure.

Why logits-level filtering?

Most LLM safety approaches work post-generation: regex on output text, classifiers on responses, or guardrails that run after the first token is produced. All of these share a fundamental weakness — the model already committed compute to generate the dangerous token.

Logits-level blocking intercepts before the token is sampled. The flow is:

  1. Vocabulary tensors arrive from the model forward pass
  2. resk-logits scans every token in the vocabulary against a prebuilt Aho-Corasick trie
  3. Matching tokens get their logits set to -inf or penalized
  4. Only safe tokens remain for sampling

This runs in a single GPU kernel launch and adds sub-millisecond overhead.

Performance

  • 10,000+ patterns mapped in under 1ms on an RTX 4090
  • 50,000+ token vocabulary scanned in one pass (not per-token)
  • CUDA and Metal backends automatically selected
  • Zero inference latency — runs before generation starts

Quick Start 

pip install resklogits
Enter fullscreen mode Exit fullscreen mode 
from resklogits import build_processor
import torch

processor = build_processor(
    patterns=["jailbreak", "ignore previous instructions"],
    penalty=-100.0,   # hard block
    device="cuda"
)

# During inference
logits = model(**inputs).logits  # [batch, seq_len, vocab_size]
safe_logits = processor(logits)  # dangerous tokens neutralized
tokens = torch.argmax(safe_logits, dim=-1)
Enter fullscreen mode Exit fullscreen mode

Features

  • Block mode — set matching token logits to -inf for hard safety
  • Shadow ban mode — apply a configurable penalty instead of hard block. The model may produce harmless alternatives without knowing censorship is active
  • Custom pattern lists — YAML or Python list input, load domain-specific blocklists (prompt injections, secret exfiltration, PII leakage)
  • Any HuggingFace model — works with any AutoModelForCausalLM that exposes .logits

Why open source?

AI security should not be a black box. With resk-logits you can audit exactly which patterns are blocked, how the trie is built, and what happens at the logits level. No magic black box, no vendor lock-in.

👉 Give it a star on GitHub if you believe LLM safety belongs at the architecture level.
👉 Check the full docs at resk.fr

Top comments (0)