Accelerating Suffix Jailbreak attacks with Prefix-Shared KV-cache

1. Problem Statement

Suffix jailbreak attacks are a critical component of "red-teaming" Large Language Models (LLMs), where adversarial suffixes are synthesized to bypass safety guardrails and elicit harmful content. However, these attacks face a severe computational bottleneck:

• Iterative Evaluation: Methods like GCG, AutoDAN, and BEAST require evaluating thousands of candidate suffixes against a target LLM to find an effective jailbreak.
• Redundant Computation: In every iteration, the "harmful instruction" (the prefix) remains static while only the suffix changes. Standard inference engines redundantly re-compute the Key-Value (KV) states for this long prefix for every single candidate suffix.
• Memory Explosion: Standard KV caching implementations (e.g., in PyTorch) require duplicating the prefix's KV cache $N$ times to match the batch size of $N$ candidates. Since the prefix is often longer than the suffix, this leads to excessive memory consumption, limiting the batch size and search width, or causing Out-Of-Memory (OOM) errors.
• Incompatibility with Advanced Systems: State-of-the-art inference engines like vLLM and SGLang optimize for multi-user serving but lack the backpropagation kernels required for gradient-based jailbreak attacks, restricting their utility to gradient-free methods only.

2. Methodology: Prefix-Shared KV Cache (PSKV)

The authors propose PSKV, a plug-and-play inference optimization framework tailored specifically for the structural redundancy in jailbreak attacks.

Core Mechanism

• Single Prefix Cache: Instead of duplicating the KV cache for the static harmful instruction ($x^{(h)}$) for every candidate, PSKV computes the KV vectors for the prefix once and stores them in a single shared cache.
• Layer-wise Broadcasting: During the inference of $K \cdot q$ candidate suffixes, PSKV retrieves the shared prefix KV cache and broadcasts (duplicates) it to the specific attention layer being processed.
• Layer-wise Expansion: Crucially, this duplication happens layer-by-layer. The system only holds the expanded KV view for the current layer in memory, discarding it before moving to the next. This prevents the simultaneous storage of redundant data across the entire network stack, drastically reducing peak memory usage.

Suffix-Centric Alignment

To handle batched attacks involving multiple distinct harmful instructions (multi-instruction scenarios), the authors introduce a Suffix-Centric Alignment strategy:

• Padding Strategy: Instructions are left-padded, and target responses are right-padded.
• Dense Tensor Formation: This ensures all generated suffixes align at the same column index, forming a dense token matrix. This allows for highly optimized vectorized operations (e.g., in PyTorch) and efficient broadcasting of the shared prefix caches across the candidate dimension.

Complexity Analysis

• Computational Complexity: PSKV maintains the same computational complexity as a standard KV cache (eliminating redundant prefix calculations) but avoids the overhead of general-purpose serving engines.
• Memory Complexity:
  • Standard KV Cache: $O(N_{cand} \cdot N_{prefix})$ (Memory scales linearly with the total number of candidates).
  • PSKV: $O(B \cdot N_{prefix} + N_{cand} \cdot N_{suffix})$ (Memory scales with the number of prompts $B$, not the total candidates). This effectively removes the memory barrier for large search widths.

3. Key Contributions

1. Identification of Redundancy: The paper identifies that general-purpose inference caching is inefficient for jailbreak attacks because it ignores the "one-prefix-many-suffixes" structural pattern.
2. PSKV Framework: Proposes a lightweight, workload-aware framework featuring virtual cache broadcasting and suffix-centric tensor alignment to eliminate redundant prefix computations and memory duplication.
3. Gradient Preservation: Unlike encapsulated inference engines (vLLM/SGLang), PSKV integrates seamlessly into existing attack pipelines (including gradient-based ones) without sacrificing access to gradients.
4. Comprehensive Evaluation: Validated across six diverse suffix attacks (GCG, GCQ, AutoDAN, BEAST, AmpleGCG, AdvPrompter) and five popular LLMs (Vicuna, Llama-2/3, Mistral, Qwen).

4. Experimental Results

The authors evaluated PSKV on 2x 80GB A100 GPUs. Key findings include:

• Performance Preservation: PSKV is "performance-neutral." The Attack Success Rate (ASR) remained statistically consistent with baseline methods (fluctuations were within the range of stochastic noise), proving the optimization does not alter the attack logic.
• Speedup:
  • Reduced inference time by ~40% on average compared to non-accelerated baselines.
  • Achieved speedup ratios of 1.4x to 1.95x depending on the attack type (e.g., AutoDAN saw nearly 2x speedup due to its high number of inference steps).
  • Outperformed or matched standard KV cache implementations without incurring time penalties.
• Memory Efficiency:
  • Reduced peak memory usage by ~50% compared to standard KV cache implementations.
  • Enabled attacks that previously resulted in OOM (Out Of Memory) errors with standard caching to run successfully.
  • Demonstrated scalability: Increasing the search width (batch size) resulted in negligible increases in peak memory, as the memory cost became dominated by intermediate tensors rather than the cache itself.
• Comparison with SOTA: While vLLM and SGLang are faster for gradient-free attacks (like BEAST), they cannot support gradient-based attacks. PSKV bridges this gap, offering comparable efficiency to these engines while supporting the full spectrum of adversarial optimization.

5. Significance

• Scalable Red-Teaming: PSKV lowers the computational barrier for evaluating LLM safety, allowing researchers to scale evaluations to larger models and broader search spaces that were previously computationally prohibitive.
• Universal Optimization: It serves as a "drop-in" optimization for the entire ecosystem of suffix jailbreak attacks, regardless of whether they are gradient-based or model-based.
• Efficiency Paradigm: The work highlights the importance of workload-specific optimizations over general-purpose inference engines for specialized security tasks, demonstrating that hard-coding structural patterns (like shared prefixes) into the inference kernel yields superior efficiency.

In conclusion, PSKV transforms suffix jailbreaking from a memory-bound, slow process into a highly efficient, scalable operation, facilitating more rigorous and comprehensive safety assessments for Large Language Models.