ml-security-research-lab/cachetrapResearch code for a gray-box trojan attack that flips a single KV-cache bit in fine-tuned LLM classifiers and measures per-class attack success rate.
Repository for "CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs", IEEE/ACM International Conference on Computer-Aided Design (ICCAD), 2026.
CacheTrap searches for a single bit flip in the KV cache of a fine-tuned LLM classifier and measures the resulting attack success rate (ASR) per target class.
Python 3.10, CUDA 12.x.
conda create -n cachetrap python=3.10
conda activate cachetrap
pip install -r requirements.txt
Hugging Face access. meta-llama/Llama-2-7b-chat-hf and
meta-llama/Llama-3.1-8B-Instruct are gated. Accept the licences on the Hub, then:
huggingface-cli login
One model / one dataset:
python train_model.py --model llama3_1_8b --dataset arc_easy
The checkpoint is written to
TrainedModels/merged_<model>_<dataset>/.
All models × all datasets:
bash scripts/train.sh
Edit MODELS, DATASETS, and NUM_GPUS at the top of the script to run a subset.
Logs go to logs/training/.
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa
--dataset is the dataset the victim was trained on; --calib_dataset is the data
available to the attacker. Following the paper, victims trained on OpenBookQA, TREC
or ARC-Challenge are calibrated with ARC-Easy, and victims trained on ARC-Easy are
calibrated with OpenBookQA. The two batch scripts encode exactly that split:
| Script | Victim datasets | Calibration |
|---|---|---|
scripts/attack_bundle1.sh | openbookqa, trec, arc_challenge | arc_easy |
scripts/attack_bundle2.sh | arc_easy | openbookqa |
bash scripts/attack_bundle1.sh
bash scripts/attack_bundle2.sh
Each ends by running summarize_logs.py over its own log directory.
Run python attack.py --help for the full option list.
logs/<run>/<model>_<dataset>.log — full per-run outputlogs/<run>/summary.csv — one row per target class (baseline accuracy, accuracy
under attack, ASR, flip location)logs/<run>/summary.txt — the raw summary blocksflip_locations/*.json — flip locations, if --save_flip_locations was passedTo re-evaluate without repeating the search:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
To try the attack without training anything, download the released
merged_llama3_1_8b_arc_easy checkpoint and place it under TrainedModels/:
TrainedModels/
└── merged_llama3_1_8b_arc_easy/
├── config.json
├── model-*.safetensors
├── tokenizer.json
└── ...
Then run the following command to attack:
python attack.py --model llama3_1_8b --dataset arc_easy --calib_dataset openbookqa \
--load_flip_locations flip_locations/llama3_1_8b_arc_easy.json
Alternatively, run the following for a fresh attack
python attack.py \
--model llama3_1_8b \
--dataset arc_easy \
--calib_dataset openbookqa \
--threat_model graybox \
--save_flip_locations flip_locations/llama3_1_8b_arc_easy_fresh.json
This prints a per-class summary of baseline accuracy, accuracy under attack, ASR, and the selected flip location.
Cite the pre-print version as
@article{nahian2025cachetrap,
title={CacheTrap: Unveiling a Stealthier Gray-Box Trojan against LLMs},
author={Nahian, Mohaiminul Al and Almalky, Abeer Matar A and Aragonda, Gamana and Zhou, Ranyang and Ahmed, Sabbir and Ponomarev, Dmitry and Yang, Li and Angizi, Shaahin and Rakin, Adnan Siraj},
journal={arXiv preprint arXiv:2511.22681},
year={2025}
}