Implementing Llama Guard: Multi-stage Safety Filtering for LLM Gateways
We run a customer-facing LLM gateway in front of a hosted model, and the threat model is unglamorous but real: prompt injection trying to exfiltrate another tenant’s data, jailbreaks attempting to bypass the system prompt, and the occasional output that should never have been generated. Rule lists and regexes are useless here — the attacks rephrase themselves faster than we can enumerate them. We needed a classifier trained on a safety taxonomy, and that’s exactly what Llama Guard is.
Llama Guard 3 (8B, released August 2024) is a Llama 3.1-based model fine-tuned to emit safe or unsafe followed by a category label. It has two properties that make it usable in a gateway: it is a small enough model to run on a single modest GPU, and its input format is deterministic, so you can build the prompt yourself and parse the first output token.
The prompt format and the two-pass design
Llama Guard 3 uses an explicit framing that the chat template handles for you: user input goes between <BEGIN USER INPUT> and <END USER INPUT>, generated content between <BEGIN CONTENT> and <END CONTENT>, and the output’s first token is the verdict. We run it twice per request: once on the prompt before it reaches the main model, once on the generated response before it is released to the caller. The second pass is non-negotiable — the first pass cannot catch a model that was already jailbroken into producing a bad response.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-Guard-3-8B", torch_dtype=torch.float16
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-Guard-3-8B")
def guard(messages):
"""messages: list of {"role": "user"|"agent", "content": str}"""
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
out = model.generate(**inputs, max_new_tokens=32, do_sample=False)
result = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
verdict = result.split()[0] if result else "unsafe"
return verdict, result # ("safe"|"unsafe", category detail)
On the gateway, the first pass is guard([{"role": "user", "content": prompt}]) and the second is guard([{"role": "user", "content": prompt}, {"role": "agent", "content": response}]) — mirroring the conversation so the classifier sees the full context, which matters because a response is only unsafe relative to what the user asked.
The latency budget, honestly
Two classification passes add real latency. On an A10G running FP8, each pass lands between 35-60ms; on an L4 with vLLM batching it is a few milliseconds slower but scales to hundreds of concurrent checks. The full request cost is roughly 120ms on top of the main model — for us that pushed p99 from ~2.1s to ~2.3s. That was the cost of the guarantees, and it is the same order of magnitude as one retry, so we accepted it.
If 120ms matters, you can shrink the exposure: run the input pass synchronously (it is the one that prevents the expensive main-model call from ever happening) and run the output pass asynchronously for low-risk flows, quarantining the response and escalating if it comes back unsafe. We do this only for a whitelisted set of low-value endpoints; everything customer-facing stays synchronous.
Shadow mode and category mapping
We shipped the whole thing in log-only mode for a week before it was allowed to block anything. The false-positive rate on real traffic was ~0.4% of legitimate prompts, almost all of it the taxonomy’s S1/S2 (violent crimes / non-violent crimes) categories catching references in product-support text — a customer asking “how do I report an attack in the simulator” is not unsafe. The 13 default categories map cleanly onto a security team’s concerns, but you must re-map them onto your domain:
- Keep default categories for PII (S13) and code-interpreter abuse — those are where the real exfiltration happens.
- Downgrade categories that fire on legitimate domain vocabulary to monitor only instead of block.
- Never silently block. Blocked traffic goes to a SIEM entry with the request ID and the exact category, and the caller gets a canned “I can’t help with that” rather than an error, so attackers can’t use the block itself as an oracle.
Production lessons
- Put Llama Guard on its own GPU pool, not the main inference pool. If the guard contends with the generator for VRAM, both latency budgets blow.
- Keep a small adversarial eval set — known jailbreak templates, injection payloads, PII probes — and gate guard-model releases on it. The threat model drifts; the eval set is what keeps the guard honest.
- Parse only the first token as the verdict and treat anything else as
unsafeby default. A truncated or malformed response must fail closed. - Guard before the call to save money as well as prevent harm: the input pass routinely stops a 30-second, token-heavy generation from ever starting.
- Shadow mode first, always. A classifier with a 0.4% false-positive rate will block real customers and produce an incident that makes leadership want to switch the thing off entirely.