Accelerating Language Model Workflows with Prompt Choreography

TL;DR

Prompt Choreography introduces a dynamic, global KV cache to accelerate multi-agent LLM workflows by allowing each call to attend to a reordered subset of previously encoded messages. The framework supports parallel decoding and efficient position management via RoPE-based rotations and on-the-fly attention masking, reducing redundant re-encoding. While the caching strategy can cause information blockage or leakage, the authors demonstrate that lightweight parameter-efficient fine-tuning (LoRA) can restore baseline or superior accuracy across representative tasks on MATH, while delivering substantial speedups (TTFT 2.0–6.2×, E2E up to 2.2×). The work also provides a reference API and implementation notes to facilitate integration into production systems, with careful attention to potential privacy and ordering effects. Overall, Prompt Choreography offers a practical path to faster, scalable LLM workflows by reusing dynamic context across agents without compromising isolation where needed.

Abstract

Large language models are increasingly deployed in multi-agent workflows. We introduce Prompt Choreography, a framework that efficiently executes LLM workflows by maintaining a dynamic, global KV cache. Each LLM call can attend to an arbitrary, reordered subset of previously encoded messages. Parallel calls are supported. Though caching messages' encodings sometimes gives different results from re-encoding them in a new context, we show in diverse settings that fine-tuning the LLM to work with the cache can help it mimic the original results. Prompt Choreography significantly reduces per-message latency (2.0--6.2$\times$ faster time-to-first-token) and achieves substantial end-to-end speedups ($>$2.2$\times$) in some workflows dominated by redundant computation.

Figures (8)

• Figure 1: Prompt Choreography manages a global KV cache of messages, which is shared and extended by all participating agents. When assembling a prompt, an agent selects messages and may specify their start positions; this allows reordering, gaps, and overlaps (not shown). Parallel decoding is possible whenever multiple prompts can be created from the current cache state.
• Figure 2: Workflows for the experiments of \ref{['sec:experiments']}. Full code sketches using our API appear in \ref{['fig:examples-full']} in \ref{['app:workflow-details']}. Each box is a message, with arrows to it from its parents. Pink boxes are prefilled. Each non-pink box has an additional parent (not shown): a system prompt corresponding to its color. For instance, each blue "Voter" in middle is generated with attention to instructions on how to select the best candidate.
• Figure 3: Attention masks used in Prompt Choreography. Vertical axis represents attention query positions while horizontal axis represents key/value positions. Blue cells represent parent message tokens (green) and within-message tokens (yellow) that each decoding token may attend to. Top (prefill) encodes 2 input messages, each consisting of 3 tokens. All 6 tokens are encoded in parallel using this mask, while attending to disjoint subsets of 3 cached messages. Bottom (decode) encodes the second tokens of 3 output messages. Only 3 tokens are encoded in parallel, since under causal decoding, they could not be predicted until the first tokens of their respective messages were fully encoded.
• Figure 4: E2E speedup of the ToT workflow across varying numbers of branches (candidate solutions) and voters. Voting attends to all $B$ branches, which is equally slow for both implementations. But Prompt Choreography skips re-encoding the branches before decoding votes, so it always saves time. Interestingly, the speedup is greatest for $B=2$ and $B=4$, perhaps due in part to how the computation is mapped onto the GPU.
• Figure 5: Attention topologies compared in MultiQA. Boxes represent messages, solid arrows represent attention dependencies, and horizontal displacement represents relative position, starting from position 0 on the left and increasing rightwards. Top depicts baseline, middle depicts choreographed serial, and bottom depicts choreographed parallel.