## Diagram: KV Cache State Transitions in a Speculative Decoding/Rollback Process ### Overview This diagram illustrates a technical process for managing Key-Value (KV) cache states during a sequence generation task, likely in the context of a language model using speculative decoding or a similar verification mechanism. The process flows from left to right through four distinct phases: Prefill, Decode, Verify, and Rollback. It shows how a "deterministic request" is processed alongside "other requests," how speculative tokens are generated and verified, and how the system state is restored upon rejection of tokens. ### Components/Axes The diagram is a flowchart with no traditional chart axes. It is organized into three horizontal sections representing the state of the KV cache at different stages, and four vertical phases detailing the process steps. **Top Section (KV Cache State):** * **Left:** "KV cache after prefill" - Represented by a single blue cylinder. * **Center:** "KV cache after decode" - Represented by a cylinder with a blue top and a yellow bottom. * **Right:** "KV cache after verify-rollback" - Represented by a cylinder with a blue top, a green middle, and a yellow bottom. **Bottom Section (Process Flow):** The process is divided into four labeled phases: 1. **Prefill** 2. **Decode** 3. **Verify** 4. **Rollback** **Key Labels and Elements:** * **Input:** "deterministic request" (pointing to a blue rectangle). * **Concurrent Work:** "other requests" (pointing to a stack of grey rectangles). * **Token Notation:** * `T0`, `T1`, `T2`, `T3`, `T4`: Likely represent actual or target tokens. * `T1'`, `T2'`, `T3'`: Represent speculative or predicted tokens (denoted with a prime symbol). * **Verification Outcomes:** * "accepted tokens" (pointing to two green checkmarks). * "rejected tokens" (pointing to two red X marks). * **Final State Description:** "Sequence and KV cache restored until the final accepted token (T2)". ### Detailed Analysis **1. Prefill Phase:** * A "deterministic request" (blue rectangle) is processed. * This occurs alongside "other requests" (grey rectangles), indicating concurrent processing. * The output is the initial "KV cache after prefill" (blue cylinder). **2. Decode Phase:** * The system generates a sequence of speculative tokens: `T0` (green), `T1'` (yellow), `T2'` (yellow), `T3'` (yellow). * These tokens are visualized as a horizontal bar. Below them are patterned blocks (grey with dots), likely representing the KV cache entries or attention patterns associated with these speculative tokens. * The "KV cache after decode" now contains the original prefill data (blue) plus new data from the speculative decode (yellow). **3. Verify Phase:** * The speculative tokens (`T0`, `T1'`, `T2'`, `T3'`) are compared against a separate sequence of actual or verified tokens (`T0`, `T1`, `T2`, `T3`, `T4`). * **Comparison Logic:** * `T1` is equal to `T1'` (`T1 = T1'`). This token is **accepted** (green checkmark). * `T2` is not equal to `T2'` (`T2 != T2'`). This token is **rejected** (red X). * The diagram implies `T0` was accepted (it matches and is the starting point). * Tokens `T3` and `T4` (red blocks) are shown but are not compared to primes, suggesting they are part of the verification sequence but were never proposed speculatively, or are beyond the point of failure. * The verification results in a set of "accepted tokens" (green checks) and "rejected tokens" (red Xs). **4. Rollback Phase:** * Based on the verification, the system rolls back. * The final state shows a sequence bar containing: the original blue block, the accepted green token (`T0`), the accepted green token (`T1`), and a patterned block for the rejected `T2`. * The "KV cache after verify-rollback" cylinder reflects this: it retains the prefill data (blue), the data from accepted tokens (green), and discards the data from rejected speculative tokens (the yellow section is now at the bottom, possibly indicating it's stale or marked for overwrite). * The text confirms: "Sequence and KV cache restored until the final accepted token (T2)". ### Key Observations * **Color Coding is Critical:** Blue represents the initial/prefill state. Green represents accepted/speculative tokens that are verified. Yellow represents speculative tokens that are generated but not yet verified. Red represents tokens in the verification sequence that are incorrect or beyond the failure point. Patterned fills represent KV cache data or states associated with tokens. * **Process Flow:** The diagram clearly shows a speculative-then-verify workflow. The system guesses multiple tokens ahead (`T1'`, `T2'`, `T3'`) in the Decode phase, then checks them in the Verify phase. * **Point of Failure:** The process fails at `T2` (`T2 != T2'`). All speculation after this point (`T3'`) is invalidated. * **State Restoration:** The Rollback phase does not revert to the very beginning. It reverts to the state after the last *accepted* token (`T2` in the sequence, which corresponds to `T1'` being correct). The KV cache is trimmed accordingly. ### Interpretation This diagram depicts a **speculative decoding with rollback** mechanism, a technique used to accelerate inference in autoregressive models like large language models. * **What it demonstrates:** The system uses a faster, possibly less accurate "draft" model to propose multiple future tokens (`T1'`, `T2'`, `T3'`) in one forward pass (Decode). A slower, more accurate "verification" model then checks these proposals against its own generated tokens (`T1`, `T2`, `T3`). If a mismatch is found (`T2 != T2'`), the system discards all subsequent speculative work and rolls back its state (KV cache and token sequence) to the last point of agreement. * **Why it matters:** This approach aims to achieve a speedup. If the draft model's guesses are often correct, the system can generate multiple tokens per verification step, reducing the total number of expensive verification passes needed. The cost of occasional rollbacks is offset by the gains from successful speculation. * **Underlying Logic:** The "deterministic request" likely refers to the initial prompt or a guaranteed starting point. The "other requests" suggest this process can be batched or run concurrently for multiple sequences. The preservation of the blue (prefill) and green (accepted) sections in the final KV cache shows efficient memory management—only the incorrect speculative data (yellow) is discarded. * **Notable Detail:** The final accepted token is labeled as `T2` in the rollback description, but in the verification column, `T1` was the last token marked with a check (`T1 = T1'`). This suggests `T2` in the final sequence bar corresponds to the *position* where the mismatch occurred, and the system restores up to and including the token *before* the failure, which is `T1`. The patterned block labeled `T2` in the rollback bar may represent the now-invalidated slot or the point from which regeneration must occur.