Skip to content

Accelerating Language Model Workflows with Prompt Choreography

🕒 Published (v1): 2025-12-28 19:21 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Prompt Choreography is a framework that replaces per-call prompt re-encoding with a shared, dynamic global KV cache across all agents in an LLM workflow. Agents select arbitrary subsets of cached messages with reordered positions via dynamic attention masking and RoPE-based position updates. A lightweight LoRA fine-tuning step mitigates accuracy degradation from modified attention semantics, yielding 2.0–6.2× faster time-to-first-token and up to 2.2× end-to-end speedup.

Problem

Standard multi-agent LLM workflows re-encode shared messages (system prompts, inter-agent outputs, background documents) redundantly on every LLM call. Prefix caching only handles static common prefixes; methods like Prompt Cache handle static precomputed documents but cannot reuse dynamically generated runtime messages. No prior approach supports arbitrary reordering of dynamically generated messages across a shared cache.

Method

A global KV cache stores (K, V) tensors for every prefilled or decoded message on the GPU. Two API primitives replace the standard complete() call: - prefill(message, parents, offsets) — encodes message tokens in parallel, attending only to specified parents at repositioned offsets; appends to cache. - decode(header, parents, offsets) — autoregressively generates a new message attending to specified parents; also appends to cache.

Position reuse is enabled by RoPE translational invariance: cached key vectors are rotated by (j' - j) to reposition a message without full recomputation. Dynamic attention masks (via FlexAttention) enforce per-call visibility rules on-the-fly. Parallel decoding interleaves tokens of concurrent messages in physical memory while maintaining logical isolation via stored (message_id, position) metadata.

When modified attention causes accuracy drops (information blockage from independently-encoded parents; information leakage from reused encodings carrying private context), a LoRA distillation step fine-tunes the model on ~100–500 baseline execution traces to recover baseline-level output distributions.

Key Contributions

  • Global dynamic KV cache API: prefill/decode primitives that decouple encoding from context assembly, supporting arbitrary message reuse, reordering, and parallel decoding.
  • RoPE-compatible position updates: efficient key-vector rotation enabling arbitrary message repositioning without re-encoding.
  • Characterization of information blockage and leakage: formal analysis and targeted experiments (MultiQA, BSM, Prisoner's Dilemma) quantifying attention-semantic side effects.
  • LoRA distillation recipe: parameter-efficient fine-tuning (<1% parameters, <3 hours on 1× A100) that recovers and sometimes exceeds baseline accuracy.
  • Reference implementation on top of Llama3.1-8B using FlexAttention, with open-source code.

Results

  • TTFT speedup (baseline ÷ choreographed): MADiter 2.0×, ToT 3.5×, MADpar 6.2× (Table 5).
  • E2E speedup: ~1.03× in standard configurations; up to 2.2× in prefill-bound ToT configurations (Figure 4).
  • MATH accuracy after fine-tuning vs. baseline: MADiter 38.6% vs. 39.0% (not significantly different); ToT 41.4% vs. 39.6% (exceeds baseline); MADpar 60.0% vs. 64.6% (small but significant gap remains) (Table 4).
  • Untuned choreography drops accuracy substantially: MADiter 24.8% (−14.2pp), MADpar 52.4% (−12.2pp).
  • MultiQA blockage recovery: fine-tuned parallel choreography restores Q1 accuracy from 32.8% → 68.1% vs. baseline 71.8% (Table 1).
  • BSM concept coverage: Choreo. Parallel + FT achieves 81.6% vs. baseline 81.0% (not significantly different) (Table 2).
  • Prisoner's Dilemma leakage: fine-tuning restores Bob's cooperation rate to within CI of baseline across all Alice strategies (Table 3).

Limitations

  • E2E speedup is modest (~1.03×) in decode-dominated workflows due to Amdahl's Law; large gains only in prefill-bound regimes.
  • Assumes the entire global KV cache fits in GPU memory (no eviction or CPU offload in the current implementation).
  • Assumes no pauses in choreography execution (e.g., human-in-the-loop or slow tool calls would hold GPU memory).
  • RoPE assumption: parallel calls sharing the same parents must place them at identical offsets; generalizing to other positional encodings requires different treatment.
  • LoRA fine-tuning does not provide strict information-isolation guarantees; leakage effects may persist at low levels.
  • MADpar accuracy after fine-tuning still trails the baseline by a statistically significant margin (60.0% vs. 64.6%).
  • Training data must be collected per workflow; a new workflow architecture requires a new distillation run.

Relevance to Harnesses / Meta-Harnesses

Prompt Choreography is directly relevant as an execution substrate for multi-agent harnesses: it provides the low-level scheduling and caching layer that a harness sits above, making the cost of spinning up parallel or iterative agent calls materially lower. A harness that fans out N agents over shared context (background documents, prior round outputs) is exactly the prefill-bound regime where Choreography achieves its largest gains. The prefill/decode API design—separating encoding from generation and exposing a programmable message-DAG—mirrors the compositional structure that meta-harnesses impose on workflow graphs, suggesting a natural integration point. The distillation recipe also offers a pattern for harness designers: workflow-level behavioral fidelity can be preserved under infrastructure changes without retraining the full model.