CodeComp: Structural KV Cache Compression for Agentic Coding

framework 2604.10235
kv-cache-compressioncode-property-graphstructured-evictioncode-generationjoernspan-protection

CodeComp: Structural KV Cache Compression for Agentic Coding — L2 #

§1 TL;DR #

CodeComp incorporates Code Property Graph (CPG) priors from Joern into KV cache compression for code tasks — span-level structural protection preserves call sites, branch predicates, and return statements that attention-only methods systematically mis-prune, recovering 91% of full-context accuracy at 60% KV retention on bug localization while adding negligible latency overhead.

§2 痛点 · 方法 · 结果 #

Q1 痛点 #

Agentic code tasks (fault localization, patch generation) process repository-scale contexts where KV cache consumes up to 70% of GPU memory. Existing compression methods (StreamingLLM, H2O, SnapKV, ChunkKV, ParallelComp) all share a common design assumption: token importance is inferred from attention scores.

This assumption breaks down for source code due to two challenges:

  1. C1 — Attention-structure mismatch: Code semantics are governed by formal program structure (control flow, data dependencies, call relationships), not token co-occurrence. Attention-based importance and CPG-derived structural importance select largely disjoint chunk sets (Jaccard overlap = 0.0944). Call sites, branch conditions, and assignments receive low attention scores despite being essential for code understanding — 52% of call sites are pruned under attention-only compression.
    1. C2 — Uniform budget allocation: Different code chunks vary greatly in structural relevance. Allocating uniform KV cache budget wastes capacity on structurally sparse regions while under-representing critical ones.
    2. Figure 1: Attention-only compression vs CodeComp on code

      Paper's Figure 1, verbatim (caption: "Attention-only compression discards structurally critical tokens such as function calls, branch predicates, and return statements. CodeComp uses CPG priors to explicitly preserve these tokens.").

      The figure illustrates the core problem: in a locate_bug function, attention-only compression prunes fetch_files(repo), f.has_error(), get_trace(f), and return trace — exactly the tokens needed for bug localization.

      Q2 方法 #

      CodeComp is a training-free, 4-step KV cache compression framework:

      1. Chunk selection: PPL-based scoring ranks code chunks by query-conditioned relevance: $s_i^{\text{ppl}} = \frac{1}{T}\sum_{t=1}^{T} -\log p(q_t \mid \text{prefix}, C_i, q_{
      2. Structural feature extraction: Joern extracts CPG for each chunk — node counts ($N_{\text{call}}$, $N_{\text{control}}$, $N_{\text{return}}$, $N_{\text{assign}}$) and edge counts ($E_{\text{cfg}}$, $E_{\text{pdg}}$).
      3. Structure-aware compression: Two complementary mechanisms:
      4. Budget allocation (C2): Distributes KV budget proportionally to structural importance — capacity multipliers in $[0.5, 1.5]$.
      5. Span protection (C1): Identifies structurally critical spans via weighted scoring (function calls 0.20, control-flow 0.18, query alignment 0.18, returns 0.14, assignments 0.14, def-use 0.10, attention 0.06). Half the chunk budget ($\rho_{\text{span}} = 0.5$) is reserved for protected spans.
      6. Attention-based residual filling: Standard attention-score aggregation fills remaining capacity after structural constraints are satisfied.
      7. 核心技术壁垒: The insight that for code, attention scores and program-structural importance are near-orthogonal (Jaccard = 0.0944), and that span-level protection dominates budget allocation in recovering accuracy. The ablation (Table 4) shows span-only achieves 0.617–0.783 accuracy while capacity-only achieves only 0.283–0.450. This means the primary bottleneck is retaining the right intra-chunk evidence, not identifying which chunk should receive more tokens.

        Q3 结果 #

        Bug localization (Table 3, Qwen3-8B, cap=0.6):

        • CodeComp avg = 0.53 — recovers 91% of full-context baseline (0.58).
        • vs ParallelComp = 0.23, SnapKV = 0.50.
        • Structural token retention (Str. Score): CodeComp 0.77 vs SnapKV 0.57 vs ParallelComp 0.52.

        Code generation (Table 2, SWE-bench Lite):

        • DS-Coder cap=0.4: CodeComp GF F1 = 0.250 vs ParallelComp 0.021 (12× improvement), patch validity 1.000.
        • Qwen-Coder cap=0.6: CodeComp GF F1 = 0.613 — 96% of uncompressed (0.637).
        • Edit distance under CodeComp matches uncompressed (0.743 vs 0.740 on DS-Coder cap=0.4).

        Ablation (Table 4):

        • Span-only (0.617–0.783) >> capacity-only (0.283–0.450) across all compression budgets.
        • Call and control-flow anchors are the most critical structural signals (Fig 4a).

        Throughput: 112–118s stable latency across retention ratios — negligible Joern overhead.

        §3 架构 / 方法图 #

        Figure 3: CodeComp pipeline overview

        Paper's Figure 3, verbatim (caption: "Overview of CodeComp, a structure-aware KV cache compression framework. Given a query and retrieved repository context, we first select relevant chunks via PPL-based scoring (Step 1). We then extract structural anchors from static program analysis using Joern and CPG (Step 2). These anchors are used to allocate chunk-level compression budgets and protect semantically critical spans (Step 3). Finally, attention-based compression fills the remaining capacity under these constraints (Step 4).").

        The 4-step pipeline: PPL chunk selection → Joern CPG extraction (call/control/return/assign nodes) → structure-aware compression (budget allocation + span protection) → constrained attention-based KV selection. Structural priors determine the critical evidence that must be preserved; attention serves as a residual selector for remaining capacity.

        Figure 2: Motivation analysis — attention-structure mismatch

        Paper's Figure 2, verbatim (caption: "Motivation analysis for KV cache compression in code. (a) Attention–structure mismatch: attention-based importance is weakly aligned with CPG-derived structural importance. (b) Structural mis-pruning: under attention-only compression, a large fraction of critical structural tokens are incorrectly discarded. (c) Structure-aware retention: incorporating structural priors significantly improves the preservation of critical structural tokens.").

        Three sub-analyses: (a) scatter plot showing weak alignment between attention and structural importance; (b) 52% of call sites, 42% of assignments pruned under attention-only compression; (c) structure-aware retention achieves 1.00 for callsite, branch, return, and signature tokens.

        Span protection mechanism #

        Within each chunk, span selection operates in two stages:

        1. Hard protection: Function signatures and query-matched spans (overlapping with query symbols) are included unconditionally.
        2. Ranked selection: Remaining spans ranked by weighted structural score, added greedily until span budget ($B_i^{\text{span}} = \min(B_i, \max(B_{\min}, \lfloor \rho_{\text{span}} B_i \rfloor))$) is exhausted.
        3. If protected tokens $|P_i| < B_i$, remaining budget filled with tokens nearest to $P_i$ in the original sequence — preserving local syntactic continuity around structurally important regions.

          System scope #

          • Stage coverage: Prefill (compression applied during prefill); decode unaffected.
          • Serving vs training: Inference serving for agentic code tasks.
          • Parallelism: Not addressed — single-model inference.
          • Deployment mode: Single node. Implemented on SGLang.

          KV / Memory manager #

          • Compression granularity: Chunk-level budget allocation + token-level span protection within chunks.
          • Compression applied: Per-layer, independently at each transformer layer.
          • Eviction criteria: Hybrid — structurally protected tokens are retained first; attention-based importance selects among remaining tokens.
          • Position encoding: Chunks retain original position indices from independent PPL scoring. Query starts at $p_{\text{query}} = \max_i \text{len}(C_i) + \text{len}(\text{prefix})$ to avoid positional conflict.

          §4 作者证明 #

          Notation table #

          SymbolMeaning
          $s_i^{\text{ppl}}$Perplexity score for chunk $i$ (lower = more relevant)
          $N_{\text{call}}, N_{\text{control}}, N_{\text{return}}, N_{\text{assign}}$CPG node counts per chunk
          $E_{\text{cfg}}, E_{\text{pdg}}$CPG edge counts (control-flow, program-dependency)
          $s(z)$Span importance score (weighted combination of structural features)
          $p^{\text{query}}(z)$Binary query-overlap indicator for span $z$
          $B_i^{\text{span}}$Span budget for chunk $i$
          $\rho_{\text{span}}$Span budget ratio (0.5)
          $u(j)$Attention-based token importance score
          $p_{\text{query}}$Starting position index for query tokens

          Equation physical meaning #

          • PPL scoring ($s_i^{\text{ppl}}$): Query-conditioned relevance — how well chunk $C_i$ predicts the query. Lower PPL = more relevant. This is a standard retrieval metric applied to chunk selection.
          • Span budget ($B_i^{\text{span}}$): Ensures minimum span protection ($B_{\min} = 16$ tokens) while capping at chunk budget. With $\rho_{\text{span}} = 0.5$, half of each chunk's capacity is structurally governed.
          • Attention aggregation ($u(j) = \sum_t A(t,j)$): Standard popularity-based importance. Used only as residual selector after structural constraints are satisfied.
          • Position encoding ($p_{\text{query}}$): Pushes query to maximum position index to avoid overlap with independently-scored chunks. Unusual RoPE handling that could interact with trained position distributions.

          6 minimum checks #

          1. ✅ Structural mismatch quantified: Jaccard = 0.0944 between top-20% attention vs structural chunks (Fig 2a).
          2. ✅ Mis-pruning quantified: 52% callsite, 42% assignment tokens pruned under attention-only (Fig 2b).
          3. ✅ Structure-aware retention validated: 1.00 for callsite/branch/return/signature (Fig 2c).
          4. ✅ Ablation isolates span vs capacity: span-only (0.617–0.783) >> capacity-only (0.283–0.450) (Table 4).
          5. ✅ Feature ablation: removing call or control-flow causes most consistent degradation (Fig 4a).
          6. ✅ Latency validated flat: 112–118s across retention ratios (Fig 4b).
          7. §5 实验与数据 #

            Setup: SGLang v0.4.9.post3. Llama-3-8B-Instruct, Qwen3-8B (bug localization); DS-Coder, Qwen-Coder (code generation). Benchmarks: InfiniteBench-CodeDebug, DebugBench, LongCodeQA, SWE-bench Lite, LCA CodeGen. Compression ratios: cap ∈ {0.4, 0.6}. Baselines: SnapKV, ParallelComp.

            Figure 4: Feature ablation and throughput

            Paper's Figure 4, verbatim (caption: "(a) Feature ablation of span scoring. (b) End-to-end latency on SWE-bench Lite.").

            Feature ablation (left) shows call and control-flow removal cause the most consistent degradation, confirming their role as primary structural signals. Latency (right) remains flat at 112–118s across all retention settings, demonstrating negligible Joern overhead in practice.

            Key results summary:

            TaskModelCapCodeCompBest BaselineImprovement
            Bug localization (DebugBench)Llama3-8B0.40.430.41 (SnapKV)+5%
            Bug localization (DebugBench)Llama3-8B0.40.430.03 (ParallelComp)14×
            Bug localization (avg)Qwen3-8B0.60.530.50 (SnapKV)+6%
            Code generation (GF F1)DS-Coder0.40.2500.200 (SnapKV)+25%
            Code generation (GF F1)DS-Coder0.40.2500.021 (ParallelComp)12×
            Code generation (GF F1)Qwen-Coder0.60.6130.700 (SnapKV)-12%
            LCA API F1DS-Coder0.60.2340.142 (ParallelComp)+65%

            Where CodeComp loses: Qwen-Coder SWE-bench at cap=0.6 — SnapKV achieves GF F1 = 0.700 vs CodeComp 0.613. SnapKV optimizes for localization F1, while CodeComp consistently achieves better generation fidelity (lower edit distance). On Qwen3-8B LongCodeQA at cap=0.4, SnapKV (0.69) > CodeComp (0.64). The two methods optimize for partially orthogonal objectives.

            Workload characterization #

            Workload regimeCodeCompAttention-only baselineWhy
            Bug localization (structure-critical)91% full-context recovery40–86% (varies wildly)Structural tokens essential for fault tracing
            Code generation (GF localization)74–96% recovery6–110% (unstable)Span protection preserves call/control tokens
            Code generation (edit distance)Matches uncompressedHigher edit distanceStructure-aware compression preserves generation fidelity
            File-level localization F1 onlySometimes weaker than SnapKVSnapKV can winAttention-based observation window better for coarse localization

            §6 论证链 #

            StepClaimEvidenceDepends on
            1Attention-based importance and program-structural importance are near-orthogonal for code§3: Jaccard = 0.0944 between top-20% chunks ranked by each signal (Fig 2a)
            2This mismatch causes systematic mis-pruning of structurally critical tokens§3: 52% callsite, 42% assignment tokens pruned (Fig 2b); ParallelComp collapses to 0.03 on DebugBench (Table 3)Step 1
            3CPG priors from Joern can identify structurally critical spans§4.2: CPG unifies AST/CFG/PDG; extracted features ($N_{\text{call}}$, $N_{\text{control}}$, etc.) align with code semantics
            4Span-level protection is the dominant accuracy mechanism§5.3: span-only (0.617–0.783) >> capacity-only (0.283–0.450); call + control-flow most critical (Fig 4a)Steps 2, 3
            5CodeComp recovers majority of full-context accuracy at reduced KV budget§5.1–5.2: 91% recovery at cap=0.6 (Table 3); matches uncompressed edit distance (Table 5)Step 4

            §7 实现 cross-reference #

            Implemented on: SGLang v0.4.9.post3. Native integration (not Transformers library).

            • CPG extraction via Joern (external tool).
            • Span scoring with configurable weights (function calls 0.20, control-flow 0.18, query alignment 0.18, return 0.14, assign 0.14, def-use 0.10, attention 0.06).
            • Chunk partitioning: function/method boundaries preferred; max 4096 tokens/chunk, target 512, min 128.
            • Attention aggregation: average pooling over sliding window (size 5), context window 128 tokens.
            • BM25 file retrieval (top-20) for SWE-bench/LCA before chunk selection.

            [实现未公开]

            关键实现细节:

            1. Span importance weights are manually tuned, not learned: The weighting (call 0.20, control 0.18, query 0.18, return 0.14, assign 0.14, def-use 0.10, attention 0.06) is fixed across all models and benchmarks. The paper does not discuss sensitivity to these weights or how they were selected — but the ablation (Fig 4a) confirms call and control-flow are the most critical, validating the relative ordering.
            2. Position encoding with overlapping chunk indices: Chunks scored independently retain their original position indices (which may overlap). The query is pushed to $p_{\text{query}} = \max_i \text{len}(C_i) + \text{len}(\text{prefix})$ to avoid positional conflict. This is an unusual RoPE handling that works empirically but could interact with the model's trained position distribution for very long repositories.
            3. API & usability #

              • User-facing API: SGLang's existing API. CodeComp is transparent to the user — compression applied internally.
              • Config surface: Retention ratio (cap), structural feature weights, span budget ratio.
              • Migration cost: Low — deploy on SGLang with CodeComp compression enabled. Requires Joern installation for CPG extraction.

              Deployment context #

              • Serving stage: Prefill (compression applied during prefill); decode unaffected.
              • Concurrency regime: Single-request evaluation; batch serving not explicitly evaluated.
              • Hardware affinity: Not hardware-specific — evaluated on standard NVIDIA GPUs.
              • Ecosystem integration: SGLang native. Joern as external dependency for CPG extraction.
              • Migration path: Install Joern + enable CodeComp compression in SGLang config. No model modification needed.