Let the Code LLM Edit Itself When You Edit the Code

## Abstract

Abstract In this work, we investigate a typical scenario in code generation where a developer edits existing code in real time and requests a code assistant, e.g., a large language model, to re-predict the next token or next line on the fly. Naively, the LLM needs to re-encode the entire KV cache to provide an accurate prediction. However, this process is computationally expensive, especially when the sequence length is long. Simply encoding the edited subsequence and integrating it to the original KV cache meets the temporal confusion problem, leading to significantly worse performance. We address this efficiency and accuracy trade-off by introducing P ositional I ntegrity E ncoding (PIE). Building upon the rotary positional encoding, PIE first removes the rotary matrices in the Key cache that introduce temporal confusion and then reapplies the correct rotary matrices. This process ensures that positional relationships between tokens are correct and requires only a single round of matrix multiplication. We validate the effectiveness of PIE through extensive experiments on the RepoBench-C-8k dataset, utilizing DeepSeek-Coder models with 1.3B, 6.7B, and 33B parameters. Our evaluation includes three real-world coding tasks: code insertion, code deletion, and multi-place code editing. Results demonstrate that PIE reduces computational overhead by over 85% compared to the standard full recomputation approach across all model sizes and tasks while well approximating the model performance.

## 1 Introduction

Large language models (LLMs) (Dettmers et al., ; Anil et al., ; Touvron et al., ; Zeng et al., ) have seen widespread adoption and achieved impressive results across various natural language processing (NLP) tasks. Despite these successes, LLMs face significant computational challenges, particularly in handling long sequences. To address this, numerous approaches have been proposed to accelerate the inference process, including lossless (e.g., memory and IO optimization (Dao et al., ; Kwon et al., ; Sheng et al., ), speculative decoding (Stern et al., ; Leviathan et al., )) and lossy techniques (e.g., quantization (Frantar et al., ; Xiao et al., ) and KV cache eviction (Xiao et al., ; Zhang et al., )). We refer to the above setting as the *static* setting, where the content is fixed, and the goal is to generate responses efficiently without compromising too much on performance.

Besides the static setting, we observe there is a strong demand for an alternative, which we call the *real-time editing* setting, where users frequently edit the content and expect the LLM to generate correct responses based on the updated information. A typical scenario is the interactive coding assistant, where developers often make incremental changes to their existing code and require the AI copilot to correctly predict the next line or complete a partial code snippet on the fly. The standard approach is re-encoding the KV cache of the content after each edit and then making the prediction. However, as illustrated in Figure 1, this approach leads to considerable substantial computational overhead and latency when the content is long  (Fu, ; Agrawal et al., ), making it impractical for real-time applications where quick and accurate responses are essential.

In this paper, we aim to improve the efficiency of AI copilots in real-time editing scenarios, as illustrated in Figure 2. Naively, an efficient strategy is encoding only the edited subsequence and then directly integrating those keys and values into the original KV cache. However, this strategy results in temporal confusion between the pre-edit and post-edit sequences. Keys of certain positions either disappear or multiple keys at different positions share the same index, causing the model to attend to incorrect information, leading to poor next-token prediction performance in practice. To address this problem, we introduce Positional Integrity Encoding (PIE). PIE is built upon rotary positional encoding (RoPE) (Su et al., ), the de-facto standard component in modern LLMs. PIE first removes the rotary matrices in the Key cache that introduce temporal confusion and then reapplies the correct rotary matrix for each position through simple matrix multiplications. By ensuring that the positional relationships between tokens in the new sequence are unique and consecutive, PIE can help the model make accurate predictions. It is worth noting that the calculation of PIE requires only a single round of matrix operations to modify the KV cache, resulting in negligible computational overhead.

We demonstrate the effectiveness of Positional Integrity Encoding (PIE) through extensive experiments conducted on the RepoBench-C-8k dataset, utilizing the DeepSeek-Coder (Guo et al., ) models with 1.3B, 6.7B, and 33B parameters. To rigorously evaluate PIE’s performance, we curated three tasks designed to simulate real-world coding scenarios: code insertion, code deletion, and multi-place code edition. These tasks were chosen to reflect common operations that developers perform during interactive coding sessions, thereby providing a comprehensive assessment of PIE’s practical utility. Our experimental results indicate that PIE achieves a reduction in computational overhead of over 85% for editing the KV cache across all model sizes and tasks without compromising performance compared to the naive full-recomputation approach. By leveraging PIE, developers can experience efficient interactions with AI coding assistants.

## 2 Related Work

#### Positional Encodings

Positional information is essential for modeling languages. The original Transformer model (Vaswani et al., ) encodes positional information using Absolute Positional Encoding (APE). In particular, a (learnable) real-valued embedding is assigned to each position $i$. Differently, Relative Positional Encodings (RPE) (Shaw et al., ; Dai et al., ; Raffel et al., ; Press et al., ; Su et al., ; Luo et al., ; ; Chi et al., ; Sun et al., ; Chi et al., ; Li et al., ; He et al., ) instead encode the relative distance $i-j$ for each position pair $(i,j)$. One of the most widely used RPE in state-of-the-art LLMs is Rotary Position Encoding (RoPE) (Su et al., ). RoPE rotates the query and key vectors by an angle proportional to their absolute positions before the attention mechanism, resulting in the attention being a function of the relative distance between tokens.

In the literature, relative positional encodings play essential roles across various tasks and data modalities, such as improving the length extrapolation capability of language models (Press et al., ; Sun et al., ; Chi et al., ; ; He et al., ) and enabling flexible modeling of structural information beyond sequence data like images (Liu et al., ) and graphs (Ying et al., ; Zhang et al., ; Luo et al., ). In this work, we develop the Positional Integrity Encoding based on RoPE to improve the efficiency of LLMs in the real-time editing setting.

#### Transformer Efficiency

Improving the efficiency of Transformer models has great significance in real-world applications. In the literature, existing approaches can be briefly categorized into (1) efficient attention, (2) model compression, and (3) system-architecture co-design.

Figure: Figure 2: Illustration of the KV cache mechanism in both static and real-time editing settings for large language models (LLMs). Top: In the static setting, the model processes a fixed input to generate predictions, leveraging precomputed Key/Value (KV) pairs stored in the cache. Bottom: In the real-time editing setting, the input code is frequently edited, necessitating updates to the KV cache to maintain accurate information to generate the correct next tokens. Our objective is to optimize the efficiency of the green arrow pathway, which represents the process of updating the KV cache in response to code edits.
Refer to caption: x2.png

The attention module of Transformer needs to calculate pairwise correlations between all positions, resulting in quadratic time and memory cost with respect to the sequence length. To reduce the cost, many efficient attention variants have been proposed, such as (1) sparse attentions (Child et al., ; Beltagy et al., ; Qiu et al., ), which design either pre-defined or learnable patterns to reduce the amount of the key-value pairs that each query needs to attend to; (2) approximation-based attention (Katharopoulos et al., ; Wang et al., ; Choromanski et al., ; Kitaev et al., ; Tay et al., ; Roy et al., ), which use tailored approaches like low-rank projection or random features to approximate standard attention for efficient computation.

Another perspective for Transformer efficiency is model compression, mainly including (1) pruning (Wang et al., ; Hubara et al., ; Ma et al., ; Frantar & Alistarh, ), which aims at removing redundant model parameters or layers for efficient deployment without scarifying performance; (2) quantization (Yao et al., ; Park et al., ; Dettmers et al., ; Frantar et al., ; Xiao et al., ; Liu et al., ), which uses post-processing to represent weights and activations via low-precision format for reducing time and memory costs; (3) knowledge distillation (Sanh et al., ; Gu et al., ), which uses a smaller model to learn knowledge from a large model for balancing efficiency-accuracy trade-offs.

In the era of LLMs, the importance of system-architecture co-design has been highlighted to improve Transformer efficiency further. Many works begin to design more efficient approaches for serving Transformers regarding the characteristics of computer systems for real-world applications, such as FlashAttention (Dao et al., ; Dao, ), PagedAttention (Kwon et al., ), and FlexGen (Sheng et al., ) that are proposed for memory and I/O optimization. Additionally, to reduce functional calls during generation, speculative decoding (Stern et al., ; Leviathan et al., ; Chen et al., ; Miao et al., ; Spector & Re, ; Cai et al., ; Zhang et al., ; He et al., ; Li et al., ) has been proposed. Our Positional Integrity Encoding is specially designed to improve the efficiency of LLMs in real-time editing settings, which can be seamlessly combined with all the above-introduced approaches to achieve further speed-ups.

## 3 Methods

### 3.1 Background

Let $s=(w_{1},w_{2},\ldots,w_{n})$ represent the input token sequence, where each $w_{i}$ belongs to a fixed vocabulary. Let $\theta_{\text{LLM}}$ represent a Transformer-based large language model, which can calculate the conditional probability distribution of the next token $p(w_{n+1}|s;\theta_{\text{LLM}})$ and generate tokens iteratively. Typically, the input $s$ is fixed. Therefore, the generation process of LLMs usually employs the KV Cache mechanism (Pope et al., ) to store previously computed Key/Value vectors during each layer’s attention calculation. We denote the KV cache as ${\bm{K}}=(K_{1},K_{2},\dots,K_{n})$ and ${\bm{V}}=(V_{1},V_{2},\dots,V_{n})$, where $K_{i}$ and $V_{i}$ are the keys and values associated with token $w_{i}$. When predicting the token at position $n+1$, we can use $w_{n}$ as input and, in each layer, compute the attention between the current hidden representation and the stored KV cache, avoiding recomputing the hidden representation of previous tokens. Without any confusion, we also denote the next token probability distribution as $p(w_{n+1}|{\bm{K}},{\bm{V}},w_{n};\theta_{\text{LLM}})$.

In this study, we aim to investigate a new scenario where the context $s$ is real-time edited by users, which makes it impossible for the KV cache to predict the correct next token without any modification. We can model such real-time context as a sequence of steps. Each step can be formulated as an action where tokens from position $i$ to $j$ in $s$ are edited, resulting in a modified sequence $s^{\text{edit}}=[w_{1},\dots,w_{i},a_{1},a_{2},\dots,a_{m},w_{j+1},\ldots,w_{n}]$, where $[a_{1},a_{2},\dots,a_{m}]$ represent the new inputs that replace $[w_{i},\dots,w_{j}]$. Our goal is to accurately and efficiently predict $w_{n+1}$ given by $\theta_{\text{LLM}}$ on $s^{\text{edit}}$. This problem is crucial in various scenarios. For instance, users can frequently edit their previous codes for different purposes and expect the code language model to swiftly adapt to these changes and predict the correct next line based on the updated information.

As the context between position $i$ and position $j$ is edited, the KV cache corresponding to these tokens must be updated. Furthermore, these changes will impact the representations of subsequent tokens after position $j$, thereby necessitating updates to all subsequent KV cache. The naive approach involves a full-recomputation strategy: re-encoding all the KV cache for $[a_{1},a_{2},\dots,a_{m},w_{j+1},\ldots,w_{n}]$ layer by layer, followed by making predictions using the updated cache ${\bm{K}}^{*}$ and ${\bm{V}}^{*}$. This approach ensures the KV cache is exact when predicting the next tokens. However, it is easy to see that it is computationally expensive, especially when the edits are light but the texts to be re-encoded are long. It’s worth noting that the original ${\bm{K}}$ and ${\bm{V}}$ already encode rich information on $[w_{j+1},\ldots,w_{n}]$, and a full recomputation may not be essential for practical problems. With this in mind, we seek to find ways to efficiently edit ${\bm{K}}$ and ${\bm{V}}$, yielding ${\bm{K}}^{\text{edit}}$ and ${\bm{V}}^{\text{edit}}$, which approximates $p(w_{n+1}|{\bm{K}}^{*},{\bm{V}}^{*},w_{n};\theta_{\text{LLM}})\approx p(w_{n+1
}|{\bm{K}}^{\text{edit}},{\bm{V}}^{\text{edit}},w_{n};\theta_{\text{LLM}})$ in an effective way.

### 3.2 Positional Integrity Encoding (PIE)

When a user modifies $s$ into $s^{\text{edit}}$, KV cache associated with the first $i$ tokens, i.e., ${\bm{K}}_{[1:i]}$ and ${\bm{V}}_{[1:i]}$, remains unchanged. As $[a_{1},a_{2},\dots,a_{m}]$ is the user’s new input, we feed this subsequence to the LLM to obtain the keys and values from position $i+1$ to $i+m$. We denote this piece of new KV cache as ${\bm{K}}^{\text{edit}}_{[i+1:i+m]}$ and ${\bm{V}}^{\text{edit}}_{[i+1:i+m]}$, and now have the edited KV cache as:

$$ $\displaystyle{\bm{K}}^{\text{edit}}$ $\displaystyle=\text{Concat}({\bm{K}}_{[1:i]},{\color[rgb]{1,0,0}\definecolor[ named]{pgfstrokecolor}{rgb}{1,0,0}{\bm{K}}^{\text{edit}}_{[i+1:i+m]}},{\bm{K}} _{[j+1:n]}),$ (1) $\displaystyle{\bm{V}}^{\text{edit}}$ $\displaystyle=\text{Concat}({\bm{V}}_{[1:i]},{\color[rgb]{1,0,0}\definecolor[ named]{pgfstrokecolor}{rgb}{1,0,0}{\bm{V}}^{\text{edit}}_{[i+1:i+m]}},{\bm{V}} _{[j+1:n]}),$ (2) $$

where the red symbols indicate real-time calculations.

#### Challenges.

The key challenge lies in how to edit the succeeding KV cache ${\bm{K}}_{[j+1:n]}$ and ${\bm{V}}_{[j+1:n]}$. Clearly, the modification of $[a_{1},a_{2},\dots,a_{m}]$ impacts the subsequent content in two ways: semantically and structurally. The semantic impact refers to the changes in the understanding of the subsequent text caused by the edited content. This can be a problem in natural language applications, such as dialog systems, where modifications to earlier conversations can significantly influence the generation of current responses. The other impact is structural, primarily concerning the temporal confusion between the pre-edit and post-edit sequences when $j-i\neq m$. This issue arises with common editing actions in code, such as additions and deletions (corresponding to $j-i=0$ or $m=0$). To be more concrete, imagine the original sequence has 5 tokens. If we add three tokens between the second and third position, it will occupy the positional index $[3,4,5]$. We calculate ${\bm{K}}^{\text{edit}}_{[3:5]}$ and ${\bm{V}}^{\text{edit}}_{[3:5]}$ in equation (1) and (2). However, in the original ${\bm{K}}_{[3:5]}$ and ${\bm{V}}_{[3:5]}$, the positional index $[3,4,5]$ is also occupied. If we integrate them together and take no actions during the next token prediction, the model will calculate similarity with multiple keys in the KV cache with the same index $[3,4,5]$, causing confusion and potential prediction errors. Empirically, in code tasks, we find that the semantic impact is relatively small. Addressing temporal confusion for light edits alone can already lead to good performance (see the experiments in Section 4 for more details).

#### Our approach.

To mitigate the temporal confusion during real-time editing, we propose a simple yet effective solution: Positional Integrity Encoding (PIE), which ensures that positional information remains correctly ordered after editing without the need to re-encode the KV cache for subsequent tokens. PIE builds upon the rotary positional encoding (RoPE) (Su et al., ), which is the most widely used positional encoding in LLMs. Without loss of generality, given a query vector ${\bm{x}}_{i}$ at position $i$ and a key vector ${\bm{x}}_{j}$ at position $j$, RoPE calculates the dot-product similarity using

$$ $\displaystyle z_{ij}={\bm{x}}^{T}_{i}W^{T}_{q}{\bm{R}}_{j-i}W_{k}{\bm{x}}_{j}$ (3) $$

where ${\bm{R}}_{j-i}$ is the rotary matrix parameterized by the relative distance $j-i$, and $W_{q}$ and $W_{k}$ are learnable projection matrices. By definition, ${\bm{R}}_{j-i}$ can be expressed by the multiplication of two rotary matrices:

$$ $\displaystyle{\bm{R}}_{j-i}={\bm{R}}_{i}^{T}{\bm{R}}_{j}$ (4) $$

For practical implementation, during inference, we compute ${\bm{R}}_{\Theta,i}W_{k}x_{i}$ as the key on the fly and store it in the cache, and when a query arrives at a new position, we rotate the query using its corresponding rotary matrix and calculate its similarity with all the keys in the cache to obtain the attention scores.

It can be easily seen that the positional information in the KV cache is encoded within the rotary matrix. When an edit occurs, the rotary matrix associated with the keys must be adjusted to reflect their post-edit locations. Leveraging the formulation of RoPE-based attention calculation, this challenge can be addressed by first removing the rotary matrices in ${\bm{K}}$ that introduce temporal confusion and then reapplying the correct rotary matrix. In detail, assume we would like to update the key vector ${\bm{k}}^{l}_{j^{\prime}}$ for the original position $j^{\prime}\in[j+1,n]$, where $l\in[1,L]$ is the layer index. We can simply edit the key vector by using

$$ $\displaystyle{\bm{k}}^{\text{edit},l}_{j^{\prime}}={\bm{R}}_{i+m+j^{\prime}-j} {\bm{R}}^{-1}_{j^{\prime}}{\bm{k}}^{l}_{j^{\prime}}$ (5) $$

where ${\bm{R}}^{-1}_{j^{\prime}}$, the inverse rotary matrix at position $j^{\prime}$, is used to remove the incorrect positional information, and ${\bm{R}}_{i+m+j^{\prime}-j}$ is used to encode the correct position $i+m+j^{\prime}-j$ in $s^{\text{edit}}$. It can easily seen that the computation can be further simplified as

$$ $\displaystyle{\bm{k}}^{\text{edit},l}_{j^{\prime}}={\bm{R}}_{i+m+j^{\prime}-j} {\bm{R}}^{-1}_{j^{\prime}}{\bm{k}}^{l}_{j^{\prime}}={\bm{R}}_{i+m+j^{\prime}-j }{\bm{R}}_{-j^{\prime}}{\bm{k}}^{l}_{j^{\prime}}={\bm{R}}_{i+m-j}{\bm{k}}^{l}_ {j^{\prime}}$ (6) $$

Hence, the full editing process for ${\bm{K}}_{[j+1:n]}$ is as follows:

$$ $\displaystyle{\bm{K}}^{\text{edit}}_{[j+1:n]}$ $\displaystyle=[K^{\text{edit}}_{j+1},\ldots,K^{\text{edit}}_{n}]$ (7) $\displaystyle\text{where each }K^{\text{edit}}_{j^{\prime}}$ $\displaystyle=\{{{{\bm{k}}}^{\text{edit},1}_{j^{\prime}},\cdots,{\bm{k}}^{ \text{edit},l}_{j^{\prime}},\cdots,{{\bm{k}}}^{\text{edit},L}_{j^{\prime}}}\}, j^{\prime}\in[j+1,n],l\in[1,L]$ $\displaystyle\text{each }{{\bm{k}}}^{\text{edit},l}_{j^{\prime}}$ $\displaystyle={\bm{R}}_{i+m-j}{\bm{k}}^{l}_{j^{\prime}}$ $$

Unlike the fully recomputation approach, the above calculation only requires a single round of matrix multiplication to directly modify the pre-computed KV cache, where the computational overhead can be considered negligible. By utilizing these transformations, we finally construct the edited KV cache as:

$$ $\displaystyle{\bm{K}}^{\text{edit}}$ $\displaystyle=\text{Concat}({\bm{K}}_{[1:i]},{\color[rgb]{1,0,0}\definecolor[ named]{pgfstrokecolor}{rgb}{1,0,0}{\bm{K}}^{\text{edit}}_{[i+1:i+m]}},{\color[ rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\bm{K}}^{\text{edit }}_{[j+1:n]})},$ (8) $\displaystyle{\bm{V}}^{\text{edit}}$ $\displaystyle=\text{Concat}({\bm{V}}_{[1:i]},{\color[rgb]{1,0,0}\definecolor[ named]{pgfstrokecolor}{rgb}{1,0,0}{\bm{V}}^{\text{edit}}_{[i+1:i+m]}},{\bm{V}} _{[j+1:n]}),$ (9) $$

where the red symbols indicate real-time calculations. The LLM then makes predictions based on $p(x_{n+1}|{\bm{K}}^{\text{edit}},{\bm{V}}^{\text{edit}},x_{n};\theta_{\text{
LLM}})$. It is worth noting that PIE is compatible with KV cache eviction methods (Xiao et al., ; Zhang et al., ; Liu et al., ). These KV cache eviction methods focus on reducing the memory usage of the KV cache during inference. PIE is designed to obtain the KV cache of the edited context with minimal overhead. By integrating PIE with KV cache eviction methods, it is possible to maintain efficient memory management while ensuring the integrity of the positional information in the real-time edit setting.

## 4 Experiments

In this section, we empirically study the effectiveness of our proposed method. In particular, we aim at answering the following questions through experiments:

- •
Question 1: Can our Positional Integrity Encoding maintain the prediction accuracy of full re-computation in code editing scenarios?
- •
Question 2: How much efficiency improvement can be achieved by using our Positional Integrity Encoding compared to existing approaches?
- •
Question 3: How large is the gap between our Positional Integrity Encoding and full re-computation in terms of LLM’s predictions & representations?

We will answer each question with carefully designed experiments in the following sub-sections. Due to space limitations, more details are presented in Appendix A.

### 4.1 Experimental Setup

**Table 1: Statistics of RepoBench-C-8k (Liu et al., ) test set.**
| Language | XF-F | XF-R | IF | Average Number of Tokens |
| --- | --- | --- | --- | --- |
| Python | 18,000 | 7,500 | 10,500 | 3,967 |
| Java | 18,000 | 7,500 | 10,500 | 4,179 |

#### Tasks.

Our experiments are conducted on RepoBench-C-8k  (Liu et al., ). This benchmark focuses on the prediction of the next line of code, given a set of in-file context (including import statements and preceding lines before the target line), and cross-file context (comprising snippets from other files parsed by import statements). The detailed statistics of RepoBench-C-8k is shown in Table 1. To effectively evaluate next-line prediction performance of code LLMs, we follow Liu et al. () to use three task settings: (1) Cross-File-First (XF-F): mask the first appearance of a cross-file line within a file; (2) Cross-File-Random (XF-R): mask a random and non-first occurrence of a cross-file line; (3) In-File (IF): mask an in-file line that does not involve any cross-file modules. Moreover, we carefully design three real-world scenarios covering code insertion, code deletion, and code edition to comprehensively examine our approach. See Appendix A for more detailed descriptions of tasks construction.

#### Settings.

In our experiments, we employ DeepSeek-Coder (Guo et al., ), a code LLM that achieves strong performance in handling repository-level code completion tasks. We use Transformers (Wolf et al., ) as our codebase. We benchmark our method on models of different sizes covering 1.3B, 6.7B, and 33B. During inference, the greedy decoding strategy is used to deterministically generate 64 tokens. For 1.3B and 6.7B models, all the experiments are conducted on a single NVIDIA A100 GPU. For 33B models, the time for encoding the context is conducted on two NVIDIA A100 GPUs and the full generation process is conducted on eight NVIDIA A100 GPUs. The first non-comment line in the output is truncated and used as the prediction. The batch size is set to 1. All experiments are repeated three times with different seeds and the averaged scores are reported.

#### Evaluation.

For comparison with our Positional Integrity Encoding, we choose two standard approaches as baselines: (1) Full-recomputation: re-compute the KV cache for all edited tokens and subsequent tokens; (2) Conflict Fast Encoding: re-compute the KV cache for the edited tokens while keeping the rest of the cache intact (i.e., using equation (1,2). Following Lu et al. (), we use Exact Match (EM) and Edit Similarity (ES) (Svyatkovskiy et al., ) to evaluate the accuracy of the predicted code lines on code completion tasks. We also report the time required to encode the edited context for efficiency evaluation.

### 4.2 Main Results

**Table 2: Performance comparisons of insertion experiments. In this task, for each next-line prediction target, we insert several lines of code into its context randomly to simulate real-world scenarios. EM and ES denotes the Exact Match and Edit Similarity score respectively. All results demonstrate that our Positional Integrity Encoding approach brings substantial speed-ups without performance drops.**
|  | Model | Method | XF-F | XF-R | IF |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | EM | ES | Time | EM | ES | Time | EM | ES | Time |  |  |
| Python | 1.3B | Full-recomputation | 22.42 | 65.26 | 192ms | 35.41 | 72.96 | 193ms | 28.78 | 69.22 | 193ms |
| 1.3B | Conflict Fast Encoding | 7.32 | 43.73 | 23ms | 10.61 | 47.18 | 22ms | 8.91 | 45.25 | 22ms |  |
| 1.3B | PIE | 22.3 | 65.2 | 29ms | 35.33 | 72.88 | 28ms | 28.69 | 69.13 | 29ms |  |
| Python | 6.7B | Full-recomputation | 28.95 | 70.11 | 561ms | 40.89 | 76.19 | 564ms | 35.26 | 72.73 | 562ms |
| 6.7B | Conflict Fast Encoding | 5.35 | 33.32 | 34ms | 6.52 | 35.25 | 34ms | 6.09 | 38.76 | 34ms |  |
| 6.7B | PIE | 28.83 | 70.01 | 50ms | 40.77 | 76.14 | 50ms | 35.2 | 72.72 | 50ms |  |
| Python | 33B | Full-recomputation | 35.75 | 73.46 | 2194ms | 46.0 | 78.9 | 2199ms | 39.75 | 75.12 | 2194ms |
| 33B | Conflict Fast Encoding | 3.96 | 30.13 | 126ms | 5.41 | 32.32 | 121ms | 3.92 | 35.56 | 127ms |  |
| 33B | PIE | 35.77 | 73.45 | 134ms | 45.74 | 78.85 | 140ms | 39.74 | 75.1 | 141ms |  |
| Java | 1.3B | Full-recomputation | 26.21 | 70.89 | 200ms | 36.77 | 76.31 | 200ms | 45.89 | 78.04 | 198ms |
| 1.3B | Conflict Fast Encoding | 0.29 | 3.12 | 22ms | 0.57 | 3.19 | 23ms | 0.7 | 2.63 | 23ms |  |
| 1.3B | PIE | 26.13 | 70.82 | 30ms | 36.67 | 76.25 | 30ms | 45.92 | 77.99 | 29ms |  |
| Java | 6.7B | Full-recomputation | 32.51 | 75.56 | 578ms | 41.97 | 79.41 | 578ms | 50.86 | 80.53 | 578ms |
| 6.7B | Conflict Fast Encoding | 0.47 | 2.77 | 34ms | 0.76 | 2.78 | 35ms | 0.67 | 2.48 | 33ms |  |
| 6.7B | PIE | 32.21 | 75.47 | 50ms | 41.96 | 79.32 | 49ms | 50.85 | 80.43 | 48ms |  |
| Java | 33B | Full-recomputation | 35.05 | 76.93 | 2269ms | 44.95 | 80.87 | 2281ms | 53.23 | 81.76 | 2270ms |
| 33B | Conflict Fast Encoding | 0.38 | 2.51 | 120ms | 0.59 | 2.4 | 122ms | 0.68 | 2.09 | 122ms |  |
| 33B | PIE | 34.78 | 76.78 | 138ms | 45.01 | 80.95 | 133ms | 53.16 | 81.66 | 139ms |  |

**Table 3: Performance comparisons of deletion experiments. In this task, for each next-line prediction target, we delete several lines of code of its context randomly to simulate real-world scenarios. EM and ES denotes the Exact Match and Edit Similarity score respectively. All results demonstrate that our Positional Integrity Encoding approach brings substantial speed-ups without performance drops.**
|  | Model | Method | XF-F | XF-R | IF |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | EM | ES | Time | EM | ES | Time | EM | ES | Time |  |  |
| Python | 1.3B | Full-recomputation | 22.42 | 65.26 | 192ms | 35.41 | 72.96 | 193ms | 28.78 | 69.22 | 193ms |
| 1.3B | Conflict Fast Encoding | 0.41 | 39.31 | 22ms | 0.59 | 42.02 | 24ms | 1.05 | 42.85 | 22ms |  |
| 1.3B | PIE | 22.31 | 65.19 | 29ms | 35.25 | 72.81 | 27ms | 28.8 | 69.16 | 27ms |  |
| Python | 6.7B | Full-recomputation | 28.95 | 70.11 | 561ms | 40.89 | 76.19 | 564ms | 35.26 | 72.73 | 562ms |
| 6.7B | Conflict Fast Encoding | 0.51 | 40.24 | 30ms | 0.77 | 42.77 | 31ms | 1.42 | 43.93 | 30ms |  |
| 6.7B | PIE | 28.86 | 70.01 | 43ms | 40.77 | 76.15 | 42ms | 35.09 | 72.67 | 42ms |  |
| Python | 33B | Full-recomputation | 35.75 | 73.46 | 2194ms | 46.0 | 78.9 | 2199ms | 39.75 | 75.12 | 2194ms |
| 33B | Conflict Fast Encoding | 1.42 | 43.93 | 105ms | 1.55 | 44.22 | 108ms | 2.4 | 45.04 | 108ms |  |
| 33B | PIE | 35.09 | 72.67 | 128ms | 45.21 | 78.18 | 118ms | 39.69 | 75.05 | 119ms |  |
| Java | 1.3B | Full-recomputation | 26.21 | 70.89 | 200ms | 36.77 | 76.31 | 200ms | 45.89 | 78.04 | 198ms |
| 1.3B | Conflict Fast Encoding | 0.33 | 33.48 | 22ms | 0.55 | 36.24 | 22ms | 1.35 | 38.65 | 22ms |  |
| 1.3B | PIE | 26.05 | 70.77 | 28ms | 36.83 | 76.34 | 27ms | 45.74 | 77.91 | 27ms |  |
| Java | 6.7B | Full-recomputation | 32.51 | 75.56 | 578ms | 41.97 | 79.41 | 578ms | 50.86 | 80.53 | 578ms |
| 6.7B | Conflict Fast Encoding | 0.49 | 33.6 | 29ms | 0.8 | 36.17 | 29ms | 1.84 | 38.82 | 29ms |  |
| 6.7B | PIE | 32.57 | 75.56 | 42ms | 41.83 | 79.39 | 43ms | 50.88 | 80.52 | 42ms |  |
| Java | 33B | Full-recomputation | 35.05 | 76.93 | 2269ms | 44.95 | 80.87 | 2281ms | 53.23 | 81.76 | 2270ms |
| 33B | Conflict Fast Encoding | 0.91 | 35.06 | 109ms | 1.01 | 37.97 | 106ms | 2.31 | 40.15 | 105ms |  |
| 33B | PIE | 34.82 | 76.76 | 117ms | 44.93 | 80.86 | 120ms | 53.19 | 81.71 | 119ms |  |

**Table 4: Performance comparisons of edition experiments. In this task, for each next-line prediction target, we delete several lines of code of its context and simultaneously insert other lines of code randomly to simulate real-world scenarios. EM and ES denotes the Exact Match and Edit Similarity score respectively. All results demonstrate that our Positional Integrity Encoding approach brings substantial speed-ups without performance drops.**
|  | Model | Method | XF-F | XF-R | IF |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | EM | ES | Time | EM | ES | Time | EM | ES | Time |  |  |
| Python | 1.3B | Full-recomputation | 22.42 | 65.26 | 242ms | 35.41 | 72.96 | 242ms | 28.78 | 69.22 | 244ms |
| 1.3B | Conflict Fast Encoding | 8.8 | 50.08 | 23ms | 13.83 | 54.01 | 22ms | 11.29 | 51.99 | 22ms |  |
| 1.3B | PIE | 22.04 | 64.59 | 30ms | 34.49 | 72.02 | 29ms | 28.15 | 68.32 | 29ms |  |
| Python | 6.7B | Full-recomputation | 28.95 | 70.11 | 705ms | 40.89 | 76.19 | 706ms | 35.26 | 72.73 | 713ms |
| 6.7B | Conflict Fast Encoding | 11.26 | 51.63 | 34ms | 12.57 | 53.27 | 34ms | 12.75 | 51.45 | 34ms |  |
| 6.7B | PIE | 28.07 | 69.0 | 54ms | 39.89 | 75.1 | 54ms | 34.05 | 71.64 | 54ms |  |
| Python | 33B | Full-recomputation | 35.75 | 73.46 | 2766ms | 46.0 | 78.9 | 2759ms | 39.75 | 75.12 | 2787ms |
| 33B | Conflict Fast Encoding | 14.33 | 53.73 | 126ms | 15.12 | 54.59 | 121ms | 13.88 | 51.94 | 127ms |  |
| 33B | PIE | 34.62 | 72.47 | 146ms | 44.59 | 77.69 | 142ms | 38.83 | 74.08 | 141ms |  |
| Java | 1.3B | Full-recomputation | 26.21 | 70.89 | 251ms | 36.77 | 76.31 | 253ms | 45.89 | 78.04 | 249ms |
| 1.3B | Conflict Fast Encoding | 5.87 | 31.46 | 22ms | 8.01 | 32.74 | 23ms | 10.2 | 34.34 | 23ms |  |
| 1.3B | PIE | 25.29 | 68.93 | 30ms | 35.59 | 74.17 | 30ms | 44.51 | 76.06 | 29ms |  |
| Java | 6.7B | Full-recomputation | 32.51 | 75.56 | 733ms | 41.97 | 79.41 | 736ms | 50.86 | 80.53 | 728ms |
| 6.7B | Conflict Fast Encoding | 9.28 | 39.32 | 34ms | 11.47 | 39.11 | 35ms | 15.51 | 41.92 | 33ms |  |
| 6.7B | PIE | 31.32 | 73.83 | 52ms | 40.89 | 77.65 | 53ms | 49.44 | 78.8 | 53ms |  |
| Java | 33B | Full-recomputation | 35.05 | 76.93 | 2892ms | 44.95 | 80.87 | 2894ms | 53.23 | 81.76 | 2833ms |
| 33B | Conflict Fast Encoding | 8.02 | 32.93 | 120ms | 9.67 | 33.29 | 122ms | 12.1 | 34.7 | 122ms |  |
| 33B | PIE | 33.72 | 74.69 | 134ms | 43.71 | 78.66 | 138ms | 51.71 | 79.49 | 143ms |  |

#### Positional Integrity Encoding perfectly preserves the full re-computation performance.

Results of different code editing settings are presented in Table 2, 3 and 4 respectively. It can be easily seen that all metrics indicate that our Positional Integrity Encoding can perfectly maintain the prediction accuracy of full re-computation in different scenarios. This indicates that PIE effectively addresses temporal confusion, and the semantic impact of minor edits is relatively negligible. For example, in the most challenging XF-F setting requiring to handle long-range cross-file context, the maximum relative difference between our PIE and Full-recomputation across different model sizes and code languages is 0.3%/0.15%, 0.66%/0.79%, 1.33%/2.24% for code insertion, deletion, and edition in terms of EM/ES respectively, which is rather negligible. This thorough examination serves as a strong support of the reliability of our PIE approach in real-world code editing scenarios.

#### Positional Integrity Encoding significantly reduces computational overhead.

Moreover, we further benchmark the computational costs brought by different approaches. From all results in the above tables, it can be easily seen that our Positional Integrity Encoding can achieve substantial speed-up compared to full re-computation while preserving its performance simultaneously. In particular, for the code edition experiment that requires both insertion and deletion, the averaged reductions of computational overhead induced by our PIE are 87.9%/88.2%, 92.4%/92.8%, 94.8%/95.2% for 1.3B, 6.7B, and 33B models on Python/Java languages respectively. Furthermore, compared to Conflict Fast Encoding which induces minimal costs but largely hurts performance, our PIE only brings negligible overhead, showing its good accuracy-efficiency balance.

In summary, our main results comprehensively demonstrate the superiority of our Positional Integration Encoding for code LLMs towards real-world code editing scenarios, which perfectly preserves the prediction accuracy and significantly addresses the crucial gap in the efficient deployment of LLMs in real-time dynamic scenarios.

### 4.3 More Analysis

In this subsection, we further present detailed analysis to investigate how large is the gap between our Positional Integrity Encoding and full re-computation in terms of context representations and predictions from LLMs, which provide additional insight of our approach.

Figure: Figure 3: Cosine similarity of key representations across model layers. The plots compare the cosine similarity between ${\bm{K}}_{[j+1:n]}$ and ${\bm{K}}^{*}_{[j+1:n]}$ (indicating temporal confusion of Conflict Fast Encoding) with the cosine similarity between ${\bm{K}}^{\text{edit}}_{[j+1:n]}$ and ${\bm{K}}^{*}_{[j+1:n]}$ (showing the effectiveness of PIE).
Refer to caption: x3.png

#### How large is the gap on context representations?

In practical scenarios, real-time editing by users results in the modified sequence $\mathbf{x}^{\text{edit}}$, requiring the KV cache to must be updated. In our analysis, we use the cosine similarity between context representations of full re-computation ${\bm{K}}^{*}_{[j+1:n]}$ and (1) our Positional Integrity Encoding ${\bm{K}}^{\text{edit}}_{[j+1:n]}$; (2) Conflict Fast Encoding ${\bm{K}}_{[j+1:n]}$. We employ the DeepSeek-Coder 6.7B model on the Python subset of RepoBench. Averaged results are reported.

In Figure 3, the cosine similarity between representations of full re-computation and our Positional Integraty Enocding is consistently around 1.0 across all layers. This high similarity demonstrates the effectiveness of PIE in preserving the contextual integrity of the key representations after editing, suggesting that PIE successfully mitigates the temporal confusion that typically arises when manipulating the KV cache. In contrast, the cosine similarity between representations of full re-computation and Conflict Fast Encoding is significantly lower, indicating the temporal confusion issue that hurts model performance a lot.

Figure: Figure 4: KL divergence of the generated token distributions. The plots compare the KL divergence between the generated token distributions of PIE and Full-recomputation, and the KL divergence between the generated token distributions of Conflict Fast Encoding and Full-recomputation.
Refer to caption: x4.png

#### How large is the gap on model predictions?

Moreover, we further use Kullback-Leibler (KL) divergence as the metric to investigate the gap between model predictions of different approaches. Similarly, we employ the DeepSeek-Coder 6.7B model on the Python subset of RepoBench and report the averaged results. In Figure 4, the KL divergence between model predictions of full re-computation and our Positional Integrity Encoding remains consistently low, i.e., below 0.0002 across 64 tokens. However, the KL divergence between model predictions of full re-computation and Conflict Fast Encoding is substantially higher (2x larger). These findings again underscore the importance of maintaining positional integrity within the KV cache to ensure accurate generation results.

## 5 Conclusion

In this paper, we introduced Positional Integrity Encoding (PIE), a novel method designed to enhance the efficiency of large language models (LLMs) in the real-time editing setting. Our approach addresses the significant computational overhead associated with re-encoding contexts after small edits, a common scenario in interactive coding environments. Through extensive experiments, we demonstrated that PIE not only significantly reduces latency but also maintains high accuracy compared to the naive full re-computation method.

PIE represents a substantial step forward in the development of efficient LLMs, particularly in dynamic contexts where frequent edits are made. Future work could explore the integration of PIE with other optimization techniques and its application to a broader range of tasks beyond code generation. Our method paves the way for more responsive and resource-efficient AI assistants, enhancing their practicality and usability in various real-world scenarios.

## Appendix A Experimental Details

### A.1 Experimental Setup

#### Tasks Construction for Code Insertion.

To simulate code insertion tasks, we start by randomly deleting five consecutive lines from each context. The resulting context, which lacks these five lines, is considered the original context. The complete context, which includes the previously deleted lines, is treated as the edited context. The tokens within the deleted lines are identified as the inserted tokens (around 64 tokens for Python and 51 tokens for Java). This setup allows us to evaluate the model’s capability to accurately restore missing code segments, mimicking real-world scenarios where developers frequently insert blocks of code.

#### Tasks Construction for Code Deletion.

For code deletion tasks, we begin by randomly selecting a line within the context and then inserting five randomly sampled lines at this position. The context containing these additional lines is designated as the original context. The complete context, which excludes the inserted lines, is regarded as the edited context. The tokens in the inserted lines are treated as the deleted tokens (around 64 tokens for Python and 51 tokens for Java). This construction enables us to assess the model’s performance in identifying and removing extraneous code, reflecting situations where developers need to clean up or refactor their codebase.

#### Tasks Construction for Multi-place Code Edition

To comprehensively evaluate the model’s performance in handling simultaneous code insertion and deletion, we construct a task scenario that integrates both operations. Initially, we randomly delete five consecutive lines from each context to simulate code insertion. The context without these lines is treated as the original context. The tokens in the deleted lines are identified as the inserted tokens.

Simultaneously, we randomly select another line within the context and insert five randomly sampled lines at this position. The complete context, which includes all lines as they appear after both deletions and insertions, is regarded as the edited context. The tokens in the newly inserted lines are considered the deleted tokens. This dual operation setup allows us to evaluate the model’s ability to handle complex, simultaneous edits, adding missing code segments while removing extraneous ones, reflecting the multifaceted nature of real-world coding environments where developers often perform multiple types of edits concurrently.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Could you provide performance results for a baseline that reuses the pre-edit context without modification for making suggestions? This zero-cost approach would be helpful to compare against Full-recomputation in your benchmark.
soundness: 1
strengths: The paper effectively outlines the real-time editing problem and clearly describes the mathematical foundation for PIE based on rotary positional encoding.
summary: This paper introduces a technique to update the rotary positional encoding when a small part of the context tokens are updated. It aims to optimize computational efficiency in real-time editing scenarios. The authors show that by fixing the positional encoding alone, they were able to retain a performance that is almost as good as fully re-encoding the context on a left-to-right code generation task, but with considerably less computational cost.
weaknesses: **Limited Technical Novelty**: The mathematical derivation is relatively straightforward, stemming directly from rotary positional encoding's relative nature, without additional innovation or complexity.

**Unrealistic Setting for Interactive Editing**:
 * Random Edits Only: The experimental setup evaluates PIE on random edits, which does not align with realistic real-time editing workflows, where temporally or contextually related edits are more common (e.g., editing a function signature and then updating its call sites). PIE’s simplifications may therefore be less applicable to typical usage patterns.
 * Single-Edit Evaluation: The paper evaluates PIE in a single-edit scenario, overlooking the potential for accumulated errors in multi-edit settings. In practical applications, users often make multiple edits, which could introduce drift in the positional encoding without full recomputation.
 * Left-to-Right Only: Evaluations are limited to left-to-right generation, omitting fill-in-the-middle (FIM) generation, a task relevant in code editing where users may modify code segments in the middle of sequences. Without this, it is unclear how PIE would perform in varied editing tasks.

**Unconvincing Conclusions Due to Limited Evaluation Scope**: Given the unrealistic evaluation settings, the claim that PIE can retain performance by adjusting positional encoding alone is unconvincing. By testing only on random edits, the experiments fail to address cases where contextual dependencies (e.g., edits that affect other tokens' relevance) might demand full recomputation. This risks overstating PIE’s applicability to real-world editing scenarios. To build a more compelling case, I recommend:
 * Evaluating PIE on real user edit sequences rather than synthetic random edits.
 * Restricting comparisons with full recomputation to cases where edited tokens impact final target tokens meaningfully (e.g., by verifying that removing or masking these edited tokens affects the target token prediction likelihood).
 * Including a special baseline where the pre-edit context is reused without modification, establishing a zero-cost apporach for comparison.

**Lack of Multi-Edit and Accumulated Error Analysis**: With each edit, additional errors will enter the encoding under the proposed technique, but the paper provides no analysis of error accumulation across multiple edits. Without such discussion, it’s unclear when a full recomputation might be needed to reset the encoding.

**Lack of Fill-in-the-Middle Evaluation**: Evaluations are limited to left-to-right generation, omitting fill-in-the-middle (FIM) generation, which is more relevant to the interactive coding assistant scenarios mentioned by the paper.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can you please add evaluations for 1 more dataset and 1 more model?
2. How does the approach do for other non-code related tasks where semantic relationship is important?
3. How does the approach do for longer edits?
4. Is there any memory overhead of the proposed approach?
5. Does this approach lead to any cumulative errors if you continue to update KV cache based on PIE?
6. The average scores of 3 experiments are reported, could you also report the standard deviation?
soundness: 4
strengths: 1. The paper is well-written and easy to understand.
2. The authors solve an important task of efficiency in updating KV cache in a real-time code editing setting. This is crucial for interactive coding assistant scenario where the developers make frequent and incremental changes to the exisiting code and require copilot to correctly predict the next line on the fly.
3. The authors perform experiments on 1 dataset for 3 tasks and show 85% reduction in computational overhead compared to brute-force approach.
summary: The paper presents Positional Integrity Encoding (PIE) as an inexpensive alternative to re-encode the entire KV cache during LLM decoding specifically for code related tasks. PIE solves the temporal confusion task efficiently using a single round of matrix-multiplicaton. The authors provide results on RepoBench dataset using 3 different sizes of DeepSeek-Coder model for 3 tasks- code insertion, deletion and multi-place editing. The results show that PIE reduces computational overhead by 85% compared to full-recomputation without significant loss of performance.
weaknesses: 1. The results are limited to 1 dataset and 1 model. Including more than 1 dataset and model would make the claim more strong.
2. The authors solve an important task of efficiency of real-time code editing but do not discuss the limitations of this approach for other tasks where semantic impact is large or in case of large code edits.
3.The approach has a dependency on RoPE and might not be suitable for other models without RoPE

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 3
questions: 1. Can the PIE be effective on code generation tasks?

2. Why not consider the Pass@1 metric?

3. Is PIE equally valid on other models (Qwen-2.5-Coder; Llama-3.1; Yi-Coder)?

I am more concerned about the third of the above issues. If the supplemental results do show the validity of PIE, I will raise my score.
soundness: 4
strengths: I really enjoy this paper!

The Positional Integrity Encoding (PIE) introduced by the authors capitalizes on RoPE, adeptly addressing temporal disorientation by initially stripping away the rotary matrices responsible for confusion and subsequently reinstating the appropriate matrices through straightforward matrix multiplication.

This capability to enhance computational efficiency without compromising accuracy is precisely the straightforward yet potent approach we value in the realm of language model optimization.

The PIE not only paves the way for future research in optimizing Large Language Models (LLMs), particularly focusing on efficiency, but also excels in real-time dynamic scenarios. Its compatibility with existing acceleration techniques positions PIE as a catalyst for further advancing the practical deployment of LLMs.
summary: This paper describes a new approach called Positional Integrity Encoding (PIE) that aims to improve the efficiency of Large Language Models in real-time code editing scenarios.

PIE improves the accuracy and efficiency of predictions by solving the problem of the computational overhead associated with recoding the entire context when editing existing code.
weaknesses: My current concerns are regarding the selection of downstream tasks and evaluation metrics considered by the authors.

(1) The tasks of code insertion, code deletion, and multi-place code editing that the authors have considered seem less critical and common in actual development scenarios compared to code generation.

(2) The chosen evaluation metrics, EM (Exact Match) and ES (Edit Similarity), may not accurately assess the semantic correctness of the generated code.

(3) The selection of models is limited to the DeepSeek-Coder series.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: - Have you conducted more experiments on more settings and models?
- Is it possible to provide an in-depth *theoretical* analysis (in 4.3) of the cause of performance drop between PIE and the recomputation?
soundness: 2
strengths: - The proposed algorithm significantly reduces the latency of generation with edited prefixes.
- A straightforward solution to adjust the rotary matrices to correct temporal confusion.
summary: This paper proposes Positional Integrity Encoding, an algorithm to correct rotary positional encoding in the KV cache of displaced tokens due to edits. It shows significant speed-up compared to full recomputation.
weaknesses: - The algorithm looks somewhat trivial to me. It simply corrects the position embedding by transforming the rotary matrix in RoPE.
- Limited experiments. All performance experiments were conducted on RepoBench-C-8k and DeepseekCoder. It is unclear if this algorithm generalizes to other settings and models.

Meta-review: This paper investigates an interesting question and presents Positional Integrity Encoding (PIE) as an inexpensive alternative to re-encoding the entire KV cache during LLM decoding, specifically for code-related tasks. The reviewers agree that the results are impressive and demonstrate significant contributions to the field. However, several suggestions for improvement remain, including issues with presentation, clarity in technical discussions, and a need for more comparisons with other contenders as well as deeper experimental analysis. The authors are encouraged to revise the paper to enhance its impact carefully.

Decision: Accept (Poster)