DEV Community

RESK
RESK

Posted on

GPU-Accelerated Shadow Ban for LLM Safety with resk-logits

#python #llm #cybersecurity #opensource 
               |         |            _) |        
  __| _ \  __| |  /      |  _ \   _` | | __|  __| 
 |    __/\__ \   <_____| | (   | (   | | |  \__ \ 
_|  \___|____/_|\_\     _|\___/ \__, |_|\__|____/ 
                                |___/
Enter fullscreen mode Exit fullscreen mode

GPU-Accelerated Shadow Ban for LLM Safety with resk-logits

Links:

The Problem

Most LLM safety filters work after generation. They regex the output, scan for keywords, or use a second model to classify the response. By the time any of that runs, the model has already generated the dangerous token. You are playing catch-up.

Prompt injection and jailbreaks exploit this gap. Attackers craft inputs that slip past instruction filters, then the model happily generates banned content because no one touched the logits.

The Approach: Logits-Level Filtering

ReskLogits operates at the logits layer, before token sampling. It uses a GPU-vectorized Aho-Corasick algorithm to build a binary mask of dangerous tokens in real time, then applies a configurable penalty to those logits.

Shadow Ban vs Hard Block

Method Mechanism Probability User Experience
Hard Block logits[token] = -inf 0% Unnatural, obvious
Shadow Ban logits[token] += -15.0 ~0.00003% Invisible

The shadow ban approach has a key advantage: the model naturally steers away from dangerous tokens without ever being explicitly blocked. This makes it harder to reverse engineer and provides a smoother user experience.

Code Example 

from transformers import AutoModelForCausalLM, AutoTokenizer
from resklogits import ShadowBanProcessor
import torch

# Load model
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Define banned phrases
banned_phrases = [
    "how to make a bomb",
    "kill yourself",
    "hack into system",
]

# Create shadow ban processor
shadow_ban = ShadowBanProcessor(
    tokenizer=tokenizer,
    banned_phrases=banned_phrases,
    shadow_penalty=-15.0,  # ~0.00003% probability
    device="cuda"
)

# Generate with protection
prompt = "Tell me how to"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
shadow_ban.reset()

outputs = model.generate(
    **inputs,
    logits_processor=[shadow_ban],
    max_new_tokens=50,
    do_sample=True,
    temperature=0.7
)

text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Generated: {text}")
# Model naturally avoids dangerous topics
Enter fullscreen mode Exit fullscreen mode

Key Features

  • GPU Vectorized Aho-Corasick: 10,000+ patterns in under 1ms on RTX 4090
  • Stateful Detection: Catches partial generations that cross token boundaries
  • Multi-Level Filtering: Configure penalties by severity tier (high/medium/low)
  • Streaming-Ready: Built-in stream() context manager and stream_generate() helper
  • Framework Compatible: Works with HuggingFace Transformers, vLLM, TGI

Multi-Level Shadow Ban 

from resklogits import MultiLevelShadowBanProcessor

phrases_by_level = {
    "high": ["bomb", "kill", "murder"],      # -20.0 penalty
    "medium": ["hack", "exploit", "crack"],   # -10.0 penalty
    "low": ["jailbreak", "bypass"]            # -5.0 penalty
}

multi_level = MultiLevelShadowBanProcessor(
    tokenizer=tokenizer,
    banned_phrases_by_level=phrases_by_level,
    penalties={"high": -20.0, "medium": -10.0, "low": -5.0}
)
Enter fullscreen mode Exit fullscreen mode

Installation 

pip install resklogits
# or with uv
uv pip install resklogits
Enter fullscreen mode Exit fullscreen mode

More Processors

ReskLogits ships with additional logits processors out of the box:

  • GenLengthLogitsProcessor - Control minimum/maximum generation length
  • BanTokenProcessor - Hard block specific token IDs or strings
  • ForceLastPhraseLogitsProcessor - Force a specific suffix at end of generation
  • MultipleChoiceLogitsProcessor - Restrict to predefined answer set
  • TriggerPhraseLogitsProcessor - Auto-response trigger on pattern match
  • CiteFromPromptLogitsProcessor - Boost tokens from the prompt (RAG)

Architecture 

[GPU] -> Logits (1 x vocab_size) -> [Vectorized Aho-Corasick] -> Mask -> Penalized Logits
Enter fullscreen mode Exit fullscreen mode

The pipeline is a single GPU kernel call. No overhead, no context switches, no latency penalty at inference time.

Try It Yourself 

pip install resklogits
Enter fullscreen mode Exit fullscreen mode

Check the full documentation on GitHub: https://github.com/Resk-Security/resk-logits
Visit resk.fr for more LLM security tools.

This article is part of the RESK Security open source series. Feedback and PRs welcome.

Top comments (0)