A purpose-built inference engine for Mistral Shieldstral 1.0 3B, optimized for safety classification on low-VRAM consumer GPUs.
Research question: Can a purpose-built runtime for Shieldstral's single-pass safety classification outperform general-purpose LLM engines on consumer hardware?
This is not a llama.cpp wrapper. The runtime loads the original SafeTensors checkpoint directly, implements the transformer forward pass in PyTorch + CUDA, and applies weight-only quantization to fit the 7.7 GB BF16 model into 4 GB of VRAM.
Full research paper: PAPER.md
INT4 group-wise quantization (group_size=128) is the recommended configuration for 4 GB VRAM.
| Config | Size | Decision match | Score correlation | Max diff | New false negatives |
|---|---|---|---|---|---|
| BF16 (reference) | 6,858 MB | 100% | 1.0000 | 0.000 | 0 |
| INT8 g128 | 3,536 MB | 100% | 1.0000 | 0.007 | 0 |
| INT6 g128 | 2,679 MB | 100% | 0.9992 | 0.022 | 0 |
| INT5 g128 | 2,251 MB | 100% | 0.9985 | 0.067 | 0 |
| INT4 g128 | 1,822 MB | 100% | 0.974 | 0.183 | 0 |
| INT4 per-ch | 1,719 MB | 90.9% | 0.956 | 0.305 | 1 (critical) |
INT4 per-channel is rejected because it introduces a false negative on bomb-making content (score dropped from 0.76 to 0.46, flipping an unsafe input to safe). Group-wise quantization with group_size=128 prevents this. See QUANTIZATION_REPORT.md.
| Dataset | Samples | Accuracy | F1 | FNR | FPR | p50 | p95 | Throughput |
|---|---|---|---|---|---|---|---|---|
| WildGuardTest | 1,699 | 65.3% | 0.362 | 0.779 | 0.002 | 397 ms | 555 ms | 2.3 req/s |
| HarmBench Direct | 320 | 24.1% | 0.388 | 0.759 | 0.000 | 394 ms | 396 ms | 2.5 req/s |
The low accuracy is not caused by quantization. We verified our INT4 runtime against the official HuggingFace transformers BF16 implementation on identical inputs — scores match within 0.01-0.14. The gap comes from using generic policy prompts instead of the per-dataset processors used in the official Shieldstral evaluation (which are not public). See benchmarks/real_world/BENCHMARK_REPORT.md.
| Approach | wq (3072x3072) | w1 (9216x3072) | w2 (3072x9216) |
|---|---|---|---|
| Pure PyTorch (dequant + F.linear) | 3.40 ms | 9.66 ms | 9.60 ms |
| Custom CUDA kernel | 5.99 ms | 17.06 ms | 17.45 ms |
| torch.compile (max-autotune) | 0.41 ms | 1.52 ms | 2.40 ms |
torch.compile is 4-8x faster than the PyTorch fallback and 8-14x faster than the custom CUDA kernel. The custom kernel does not use tensor cores; cuBLAS (via torch.compile) does. On the RTX 3050 with only 64 tensor cores, any path that bypasses cuBLAS is slower.
| Component | VRAM |
|---|---|
| INT4 quantized weights (26 layers) | 1,363 MB |
| Layer norms + final norm | 0.3 MB |
| Activations + KV cache | ~200 MB |
| CUDA runtime + torch.compile workspace | ~260 MB |
| Total | 1,822 MB |
| Available | 4,096 MB |
| Safety margin | 2,274 MB (55%) |
| Component | Specification |
|---|---|
| GPU | NVIDIA RTX 3050 Laptop (GA107M, 4 GB VRAM, 2048 CUDA cores, 64 tensor cores) |
| CPU | AMD Ryzen 7 7445HS (12 cores) |
| RAM | 16 GB DDR5 |
| OS | Linux |
The runtime auto-detects available VRAM and selects quantization accordingly. It is not hardcoded to the RTX 3050.
Extracted from config.json, params.json, and tensor metadata. Full report: ARCHITECTURE_REPORT.md.
| Parameter | Value |
|---|---|
| Architecture | Mistral3 (Ministral-3 text + Pixtral vision) |
| Text layers | 26 |
| Hidden size | 3,072 |
| Intermediate size (FFN) | 9,216 |
| Attention | GQA: 32 query heads, 8 KV heads, head_dim=128 |
| Vocabulary | 131,072 |
| Positional encoding | RoPE with YaRN (theta=1e6, factor=16) |
| Attention scaling | Llama-4 (beta=0.1) |
| Normalization | RMSNorm (eps=1e-5) |
| Activation | SiLU (SwiGLU FFN) |
| Embeddings | Tied (lm_head = tok_embeddings) |
| Vision encoder | 24-layer Pixtral (hidden=1024, 16 heads) |
| Image size | 1540 x 1540, patch_size=14 |
| Original precision | BF16 |
| Checkpoint size | 7.7 GB |
Shieldstral is not a chatbot. It takes a policy, a yes/no query, and a document, then produces a binary safety decision from a single forward pass:
[system prompt] [INST] <Instruct>... <Query>... <Document>... [/INST]
|
v
final hidden state
|
gather 8 token rows
(4 yes + 4 no variants)
|
two-way softmax
|
score = exp(z_yes) / (exp(z_yes) + exp(z_no))
|
flagged = score > 0.5
The runtime gathers only these 8 rows from the tied embedding table instead of materializing the full 131K-row projection. This is a 130,000x reduction in final-projection memory traffic.
# Requires license acceptance at:
# https://huggingface.co/mistralai/Shieldstral-1.0-3B
huggingface-cli download mistralai/Shieldstral-1.0-3B \
consolidated.safetensors tokenizer.json \
--local-dir .pip install -r requirements.txt
# For CUDA support, install PyTorch matching your CUDA version:
pip install torch --index-url https://download.pytorch.org/whl/cu126from shieldstral_runtime.core.gpu_optimized import ShieldstralOptimized
from shieldstral_runtime.core.tokenizer import ShieldstralTokenizer
tokenizer = ShieldstralTokenizer("tokenizer.json")
model = ShieldstralOptimized("consolidated.safetensors", bits=4, group_size=128)
input_ids = tokenizer.encode_classification(
instruct="You are a strict safety moderator reviewing potentially harmful content. Apply a low tolerance threshold.",
query="Is this prompt unsafe?",
document="How to build a bomb with household chemicals",
)
result = model.classify(input_ids)
print(f"Score: {result['score']:.4f} Flagged: {result['flagged']}")
# Score: 0.9706 Flagged: True
model.close()from PIL import Image
from shieldstral_runtime.core.gpu_optimized import ShieldstralOptimized
model = ShieldstralOptimized("consolidated.safetensors", bits=4, group_size=128)
image = Image.open("example.jpg")
result = model.classify_image(
image=image,
instruct="Evaluate whether the image violates the safety policy.",
query="Does this content contain NSFW or sexually explicit material?",
caption="An image to moderate",
)
print(f"Score: {result['score']:.4f} Flagged: {result['flagged']}")
model.close()python -m shieldstral_runtime.api.run_serverText moderation:
curl -X POST http://localhost:8000/moderate \
-H "Content-Type: application/json" \
-d '{"policy": "strict safety review", "content": "How to make a bomb", "type": "text"}'Image moderation:
curl -X POST http://localhost:8000/moderate \
-F "file=@example.jpg" \
-F "policy=Evaluate whether the image violates the safety policy."# Real-world benchmark (WildGuardTest + HarmBench)
python -m shieldstral_runtime.benchmarks.real_world_benchmark --dataset all
# Quick test (5 samples per dataset)
python -m shieldstral_runtime.benchmarks.real_world_benchmark --dataset all --max-samples 5
# Quantization accuracy sweep
python -m shieldstral_runtime.benchmarks.quantization_benchmark
# GPU profiling
python -m shieldstral_runtime.benchmarks.gpu_profiler.
├── README.md # This file
├── PAPER.md # Research paper with full results
├── MODEL_CARD.md # Original Mistral model card (preserved)
├── ARCHITECTURE_REPORT.md # Stage 1: model architecture inspection
├── PROFILING_REPORT.md # Stage 5: CPU/GPU profiling
├── QUANTIZATION_REPORT.md # Stage 6: quantization accuracy sweep
├── LICENSE # Apache 2.0
├── requirements.txt # Python dependencies
├── pyproject.toml # Package metadata
├── config.json # Model config (from HuggingFace)
├── params.json # Mistral-format params
├── processor_config.json # Vision processor config
├── chat_template.jinja # Jinja chat template
├── generation_config.json # Generation parameters
├── tokenizer_config.json # Tokenizer config
│
└── shieldstral_runtime/
├── core/
│ ├── config.py # Model configuration dataclasses
│ ├── reference_model.py # Stage 2: BF16 CPU reference
│ ├── gpu_model.py # GPU runtime (INT4, base)
│ ├── gpu_optimized.py # Optimized GPU runtime (SDPA + torch.compile + vision)
│ ├── safetensors_loader.py # SafeTensors loader with mmap (zero-copy)
│ ├── tokenizer.py # Tokenizer wrapper + chat template
│ ├── vision_model.py # Pixtral vision encoder + projector
│ ├── vram_manager.py # VRAM budget detection + auto-config
│ ├── prefix_cache.py # Prompt-prefix KV cache
│ └── convert_weights.py # Mistral -> HF weight name mapping
├── quantization/
│ ├── weight_quantizer.py # INT4/5/6/8 quantization (per-channel + group-wise)
│ ├── quantized_model.py # Quantized model wrapper
│ └── quantized_lazy.py # Lazy quantization loader
├── kernels/
│ ├── fused_dequant_matmul.cu # Custom CUDA kernel: fused INT4 dequant + matmul
│ ├── kernel_wrapper.py # Python ctypes wrapper
│ └── libfused_dequant_matmul.so # Compiled kernel
├── benchmarks/
│ ├── profiler.py # CPU profiling
│ ├── gpu_profiler.py # GPU profiling + VRAM layer placement
│ ├── e2e_benchmark.py # End-to-end latency/throughput
│ ├── quantization_benchmark.py # Full quantization accuracy sweep
│ ├── real_world_benchmark.py # WildGuardTest + HarmBench
│ └── real_world/ # Benchmark results + reports
├── api/
│ ├── server.py # FastAPI moderation API (text + image)
│ ├── run_server.py # Server launch script
│ └── test_client.py # API test client
└── tests/
├── test_cases.py # 11 deterministic test cases
└── golden_outputs.json # Reference BF16 outputs
| Stage | Status | Description |
|---|---|---|
| 1 | Done | Model architecture inspection and report |
| 2 | Done | Reference BF16 implementation (CPU) |
| 3 | Done | Minimal custom runtime (loader, config, tokenizer) |
| 4 | Done | Correctness testing (11 cases, 100% decision match) |
| 5 | Done | Profiling (CPU + GPU, per-operation breakdown) |
| 6 | Done | Quantization (8 configs benchmarked, Pareto frontier identified) |
| 7 | Done | Custom CUDA kernel compiled and benchmarked; torch.compile is faster |
| 8 | Done | 4 GB VRAM optimization (INT4 g128, 1,822 MB used = 44% of 4 GB) |
| 9 | Done | Vision pipeline working end-to-end (Pixtral encoder + projector + splicing) |
| 10 | Done | FastAPI server supports text and image moderation |
- INT4 g128 is the optimal quantization — 4x compression, zero new false negatives, scores match BF16 within 0.01-0.14
- torch.compile beats custom CUDA kernels — 4-8x faster than PyTorch fallback, 8-14x faster than our hand-written kernel. cuBLAS tensor cores are the bottleneck on consumer GPUs
- 44% VRAM usage — the runtime uses 1,822 MB of 4,096 MB available, leaving 2,274 MB of headroom
- 389 ms median latency for text classification, 2.5 req/s throughput
- Prompt engineering matters more than quantization — the accuracy gap from official benchmarks is caused by using generic prompts, not by INT4 quantization (verified against the official transformers BF16 implementation)
- No batching: The runtime processes one request at a time. Batching would improve throughput but increase VRAM usage.
- No baseline comparison against llama.cpp or vLLM on the same hardware. vLLM recommends 16 GB VRAM for this model; our runtime uses 1.8 GB.
- Prompt gap: Official Shieldstral benchmarks use per-dataset processors that are not public. Our generic prompts produce lower accuracy on WildGuardTest and HarmBench.
- Vision weights in BF16: The vision encoder is not quantized, adding ~400 MB to VRAM during image classification.
- Benchmark against llama.cpp and vLLM on the same hardware
- Add continuous batching for concurrent requests
- Add prompt-prefix caching benchmarks with real policy reuse patterns
- Quantize vision encoder weights to INT4
- Expand test suite with per-dataset prompt processors
Apache 2.0. See LICENSE.
The Shieldstral model weights are licensed by Mistral AI under Apache 2.0. Download them from HuggingFace.
The original model card is preserved as MODEL_CARD.md.