Lengthy Context Isn’t Free — I Constructed a Protected Immediate-Pruning Layer That Makes LLM Techniques Work

I’ve labored on, dialog state tends to develop shortly over time. It’s frequent to resend massive parts of the historical past on every flip—together with older device outputs, repeated RAG retrievals, and context that’s not related. As this accumulates, prompts can turn out to be considerably bigger, which can improve inference price and latency, and in some circumstances have an effect on reasoning efficiency.

I constructed a deterministic pipeline that prunes this redundant state earlier than the immediate ever reaches the mannequin. The model I carried out avoids LLM calls, embeddings, and exterior dependencies. Relying strictly on normal library parts ensures each pruning resolution stays totally deterministic and reproducible.

State monitoring executes in three distinct passes: Expired Context Elimination, Duplicate Context Elimination, and Dependency Restoration. The third go is what helps make the primary two safer in apply. It ensures that nothing a later message depends upon is by chance eliminated.

Whereas constructing it, I bumped into two bugs that modified the design. My first benchmark corpus used a set variety of duplicates and off device calls, which made discount percentages shrink as conversations grew. That didn’t mirror what I’d anticipate from real-world conduct.

My dependency restoration logic additionally went fully untested at first as a result of my artificial knowledge by no means created a case the place a required message was really eliminated. Each points are lined right here, together with how fixing them modified the outcomes.

After correcting the pipeline, I benchmarked it throughout three workloads: plain chat, a RAG assistant, and a tool-heavy agent. Every was examined at 5 dialog sizes, for a complete of 15 configurations on two completely different machines. Throughout these runs, all labeled required information had been preserved. The system additionally reached a secure mounted level after a single go, which suggests pruning an already pruned immediate produces no additional adjustments.

Token discount depends upon the workload. It’s about 2 to 4 p.c for plain chat, 27 to 32 p.c for a RAG assistant, and 33 to 34 p.c for a tool-heavy agent. Even at 2,000 turns and 131,000 tokens, preprocessing stayed beneath 50 milliseconds.

Full code, all 35 assessments, and the uncooked terminal output are included beneath so you’ll be able to run the pipeline your self.

Full code: https://github.com/Emmimal/prompt-pruning-layer/

Each immediate I ship retains getting heavier

I stored seeing the very same sample pop up throughout each long-running agent I constructed. A dialog begins out completely clear. Fifty turns later, it’s an absolute mess.

By that fiftieth flip, the immediate payload you’re transport out on each single request consists of the system immediate, the complete chat historical past, 4 device outputs (two of that are fully stale as a result of the device ran twice), six retrieved chunks (three are near-duplicates as a result of the person circled again to an previous subject), a SQL end result from twenty turns in the past that’s simply sitting there, and a single person desire acknowledged as soon as and by no means introduced up once more.

None of these items is technically fallacious. Each single piece made sense the precise second it was injected. The actual problem is that nothing ever will get cleared out. The immediate turns into an append-only log of each historic occasion, and also you’re dumping the entire thing onto the mannequin, each single flip, indefinitely.

That’s not only a storage drawback. Reasoning efficiency measurably degrades as enter size grows, even when the added content material is irrelevant to the duty [1], and fashions are worse at utilizing data buried in the midst of a protracted context than data close to the perimeters [2].

The speedy knee-jerk response if you see this bloating is simply primary truncation. Slice the final N messages and drop all the things else. It’s actually one line of code. The catch is it silently breaks your dialog chains in methods you gained’t even notice till it blows up on an actual person:

Flip 3:   Person: My most popular output format is CSV.
Flip 4:   Assistant: Received it.
...
Flip 47:  Person: Export the outcomes.

In case your window is capped on the final 20 turns, flip 3 disappears by flip 47. The assistant loses the context that the person requested for a CSV. That knowledge was not stale; it was a tough dependency, and easy positional truncation can’t differentiate between previous, expendable context and previous context {that a} later flip nonetheless depends on.

That’s the precise design constraint this venture addresses. Any mechanism that removes redundant state should distinguish between these two classes. Recency alone is inadequate.

Borrowing an thought from working methods

Right here is the reframe I constructed the remainder of this round: an working system is continually deciding which pages keep resident in RAM and which get evicted. An extended-running LLM dialog has the very same drawback, besides nothing performs the function of the reminiscence supervisor. Context simply accumulates without end as a result of no course of owns the job of deciding what stops incomes its place within the immediate.

The Immediate Pruner on this article is that lacking piece.

Each piece of dialog state (a person flip, an assistant flip, a device output, a retrieved chunk) is a Message object with a task, a flip quantity, and a little bit of bookkeeping metadata:

@dataclass
class Message:
    id: str
    function: str
    content material: str
    flip: int
    tool_call_key: Non-obligatory[str] = None
    expires_after_turn: Non-obligatory[int] = None
    defines_keys: listing = subject(default_factory=listing)

Three passes run over that listing so as. Every one is a pure perform: messages in, a filtered listing out. No mannequin is concerned anyplace within the pruning step itself.

Why there’s no mannequin contained in the pruner itself

I may have used an embedding mannequin to attain message relevance, which might have made this fuzzier and rather a lot simpler to write down. I didn’t, and it isn’t for the sake of purity.

As soon as a pruning resolution depends upon a mannequin’s judgment, you lose the flexibility to purpose about what a dialog will appear like on the following flip. The identical enter not ensures the identical output. A pruning layer meant to make a manufacturing system extra predictable shouldn’t be the least predictable a part of it. Every thing right here runs on dataclasses, regex, and dict lookups, which is precisely the toolset this drawback wants.

Go 1: Expired Context Elimination

If a device will get known as greater than as soon as beneath the identical key (the identical search question, the identical SQL lookup, the identical file learn), solely the latest end result continues to be reliable. Every thing earlier beneath that secret’s expired.

def _pass1_expired_context_elimination(self, messages):
    last_occurrence = {}
    for m in messages:
        if m.tool_call_key:
            last_occurrence[m.tool_call_key] = m.id

    stored, eliminated = [], []
    for m in messages:
        if m.tool_call_key and last_occurrence[m.tool_call_key] != m.id:
            eliminated.append(m)
        else:
            stored.append(m)
    return stored, eliminated

Go 2: Duplicate Context Elimination

Retrieval pipelines continuously pull up similar or near-duplicate passages, particularly when a person loops again to an earlier subject. This go normalizes the whitespace and casing, retains solely the primary prevalence, and drops each duplicate after it.

def _pass2_duplicate_context_elimination(self, messages):
    seen, stored, eliminated = {}, [], []
    for m in messages:
        if m.function == ROLE_RETRIEVED_DOC:
            norm = " ".be part of(m.content material.decrease().cut up())
            if norm in seen:
                eliminated.append(m)
                proceed
            seen[norm] = m.id
        stored.append(m)
    return stored, eliminated

Go 3: Dependency Restoration (and the bug that made me construct it correctly)

This go is the rationale the primary two are protected to run in any respect. If Go 1 or Go 2 drops a message that occurs to be the one place a still-referenced truth bought outlined, this go catches it and places it again.

The mechanism is deliberately easy: a message marks itself with a literal DEFINE: tag, and a later message references it utilizing REF:. If a REF survives into the ultimate stored set however its matching DEFINE bought dropped upstream, Dependency Restoration restores that lacking message to the listing.

def _pass3_dependency_restoration(self, all_messages, kept_messages, removed_messages):
    kept_ids = {m.id for m in kept_messages}
    by_id = {m.id: m for m in all_messages}

    key_definer = {}
    for m in all_messages:
        for key in m.defines_keys:
            key_definer[key] = m.id

    referenced_keys = set()
    for m in kept_messages:
        referenced_keys.replace(m.references())

    restored = []
    for key in referenced_keys:
        definer_id = key_definer.get(key)
        if definer_id and definer_id not in kept_ids:
            restored_msg = by_id[definer_id]
            kept_messages.append(restored_msg)
            kept_ids.add(definer_id)
            restored.append(restored_msg)

    kept_messages.type(key=lambda m: (m.flip, m.id))
    return kept_messages, restored

Right here is the bug. I ran my first full benchmark throughout all three workloads and 5 sizes, and each single row printed “Restored (deps): 0.” Each one. My first response was that this regarded nice—an ideal security document. It wasn’t. It meant Go 3 had by no means really restored a single message on any run at any measurement.

I went again into my corpus generator to search out out why, and the reply was embarrassing as soon as I noticed it. My artificial conversations solely ever connected DEFINE markers to plain person messages, whereas Go 1 and Go 2 solely ever take away device outputs and retrieved paperwork. The 2 classes by no means overlapped. Dependency Restoration was sitting within the codebase, totally written, however fully untested by my very own benchmark as a result of nothing I generated ever gave it a purpose to fireside.

The repair was to let some device outputs additionally outline a dependency, the identical method an actual “get person settings” device name may floor a truth the dialog depends upon later, although that precise device name may get outmoded and marked expired by Go 1. As soon as I added that, the numbers modified instantly: 2 restorations on the smallest tool-agent dialog, climbing to 127 on the largest. Required information stayed at one hundred pc preserved the complete time, however now that quantity really meant one thing, as a result of the go being examined was able to failing and didn’t.

I need to be direct concerning the limitation that’s nonetheless right here even after the repair: dependency detection is literal identifier matching, not understanding. It catches an actual REF matching an actual DEFINE. It won’t catch a person paraphrasing, asking “what format did I point out earlier” with no matching tag. A semantic dependency resolver would wish an embedding mannequin or an LLM name, and that’s explicitly outdoors what this deterministic pipeline does. I’d somewhat ship a narrower assure I can show than a broader one I can’t.

Design targets, and why every one exists

Deterministic. Identical enter all the time yields the identical output. Maintaining the mannequin out of the loop eliminates run-to-run variance and any threat of hallucinating what ought to stay within the immediate.

Dependency-safe. It by no means silently drops a truth {that a} later flip nonetheless wants. Positional truncation fully lacks this property, which makes it non-negotiable right here. A pruner saving 40 p.c of tokens that sometimes breaks a dialog is a worse trade-off than one saving 4 p.c that by no means breaks something.

Idempotent. Working the pruner a second time does nothing. If that’s not true, you can not safely re-prune on each single flip of a rising dialog with out risking compounding drift.

Light-weight. The pruning step ought to by no means turn out to be the precise bottleneck it was constructed to remove.

The benchmark: three workloads, and a mistake I virtually shipped

My first model of the artificial corpus generator picked a set variety of duplicate passages and repeated device calls (six device calls and eight duplicates) no matter how lengthy the dialog was. I ran it at 5 sizes and the discount proportion went down because the dialog bought longer: 9.9 p.c at 50 turns, dropping to 0.3 p.c at 2,000 turns. That runs backward from precise manufacturing visitors, and it isn’t defensible. If the quantity of waste in a benchmark is only a fixed I hand-picked, anybody studying it’s proper to ask whether or not I constructed the benchmark to show the algorithm works.

So I threw that generator out and rebuilt it round an express workload mannequin as an alternative: retrieval-per-turn, retrieval overlap fee, device name fee, and gear repetition fee. All of those had been mounted earlier than working a single benchmark, primarily based on what appeared like believable manufacturing conduct, somewhat than adjusted afterward to hit a quantity I preferred. Three workloads got here out of that:

Regular chat. No retrieval, occasional device calls, largely peculiar back-and-forth.

RAG assistant. Retrieves paperwork on each flip, with an actual probability any given passage overlaps one thing retrieved just lately, as a result of customers revisit subjects and retrieval re-surfaces the identical chunks.

Instrument agent. Frequent calls throughout 5 device varieties (search, SQL, calculator, filesystem, net fetch), excessive repetition fee, modeling one thing that re-plans and re-queries continuously.

Each artificial corpus additionally ships with floor fact: each message some later message depends upon will get labeled required up entrance. So “did pruning preserve all the things it wanted to” is a examine in opposition to identified labels, not a guess.

Right here is the entire output. All 15 configurations, not a slice of it:

Each single row says Details stored: Sure and Idempotent: Sure. Not most rows. All fifteen.

The sample by workload is sensible when you have a look at the place the waste really comes from. Regular chat barely retrieves something and barely repeats a device name, so there’s virtually nothing for Go 1 or Go 2 to catch; it stays round 4 p.c regardless of how lengthy the dialog runs. The RAG assistant retrieves each flip with actual overlap, so Duplicate Context Elimination carries a lot of the weight, touchdown round 32 p.c. The Instrument agent combines each issues (frequent device repetition and retrieval overlap) and hits the best discount at 33 to 34 p.c.

Completely different workloads accumulate completely different sorts of waste. The pruner responds on to no matter waste is sitting in entrance of it, somewhat than producing a flat quantity that might counsel the benchmark was reverse-engineered to hit a goal.

If I solely bought to maintain one end result from this entire benchmark, it’s the security property: 15 out of 15 configurations preserved one hundred pc of required information. Zero lacking dependencies, throughout three structurally completely different workloads and a 40x vary in dialog size. Truncation by place can’t supply that. For me, this was a very powerful sign when evaluating whether or not the strategy could be usable in manufacturing.

That quantity can be computed, not hand-counted. The benchmark script itself tallies how lots of the 15 (workload, measurement) pairs preserved each required truth and what number of reached the idempotent mounted level, printing an mixture abstract block after the 15 detailed runs:

============================================================
SUMMARY
============================================================
Configurations run:               15 (3 workloads x 5 sizes)
Required information preserved:         15/15
Reached mounted level (idempotent): 15/15

Workload            Token discount vary    Details preserved   Idempotent
Regular chat          1.9-4.1%                YES               YES
RAG assistant        27.0-32.5%               YES               YES
Instrument agent           33.1-33.7%               YES               YES

I wrote this examine as a result of I caught myself hand-counting desk rows for an early draft. That metric belongs within the script output, not my very own eyes squinting at a terminal log. If a take a look at run ever hits something lower than 15/15, it means I broke the pruner and have a regression to search out, not a typo to edit within the put up.

I left one particular metric out of the ultimate numbers: tokens eliminated per millisecond of execution overhead. The code computes it—it peaks at round 4,000 tokens per millisecond on small tool-agent runs and lands between 900 and 1,700 tokens at bigger scales. It stays within the codebase as an inner subject as a result of it helps observe scaling prices, nevertheless it belongs outdoors the primary desk. Readers can’t act on it the way in which they’ll with uncooked token rely, discount proportion, or millisecond overhead. Three direct metrics displaying a transparent trade-off are higher than a fourth that acts as a novelty.

The idempotence result’s the half I preferred monitoring essentially the most. Proving prune(prune(x)) == prune(x) means the pipeline hits a secure mounted level on the primary go. Working it once more on an already-pruned immediate adjustments nothing:

That guidelines out oscillation. It additionally guidelines out cumulative shrinkage throughout turns in the event you re-pruner on each single message of a rising dialog, which is precisely how this runs in manufacturing.

Reproducing it on a second machine

I ran the complete benchmark on a Linux container working Python 3.12.3, then once more on Home windows 11 in PyCharm utilizing Python 3.12 in a separate venv. Each token rely and message rely matched precisely throughout each machines. Solely the millisecond timings moved, which is normal for various {hardware}, and even these stayed beneath 50 milliseconds for the most important, most tool-heavy dialog on each setups.

One factor I observed whereas evaluating the 2 runs and need to be straight about: prompt-build time (the time to serialize the ultimate message listing right into a string) sometimes got here out slower after pruning than earlier than on the Regular Chat workload. Not by a lot—beneath a millisecond—however the course was backward.

My learn is that Regular Chat solely removes 2 to 4 p.c of messages, so the earlier than and after listing sizes are practically similar. At sub-two-millisecond operations, system jitter and rubbish assortment pauses simply swamp the precise sign. Including a warm-up name and utilizing the median of 30 runs largely stabilized the metric, however the anomaly nonetheless pops up on that workload. It’s noise at a scale the place the delta is smaller than the measurement error, so I left it uncooked somewhat than massaging the information.

What this benchmark intentionally doesn’t measure

This benchmark isolates token discount and pruning overhead as a result of these are the metrics the pipeline really controls. Finish-to-end LLM latency is a very separate variable. It depends upon supplier structure, batching, regional caching, and community situations that this venture can’t see. Making an attempt to transform a token discount proportion straight right into a latency delta means inventing an arbitrary conversion fixed.

The baseline actuality is straightforward: chopping 30 to 34 p.c of enter tokens means the mannequin does much less work per name. Usually, inference price and latency have a tendency to extend with immediate measurement [4], making this a helpful price lever. However a real latency quantity requires a stay validation go in opposition to your particular supplier. Publishing a generic latency determine right here would imply making claims about infrastructure I don’t management, somewhat than evaluating the pruning layer itself.

How this matches into an actual agent loop

The pruner lives in precisely one spot: proper after the dialog historical past is pulled collectively for a flip, and simply earlier than it will get serialized into the ultimate immediate string.

from prompt_pruning import PromptPruner, PromptBuilder

pruner = PromptPruner()
builder = PromptBuilder()

def handle_turn(conversation_state, new_user_message):
    conversation_state.append(new_user_message)

    pruned_messages, report = pruner.prune(conversation_state)
    immediate = builder.construct(pruned_messages)
    response = call_llm(immediate)

    conversation_state.append(make_assistant_message(response))
    return response

As a result of the pipeline is idempotent, calling prune() each flip of a rising dialog is protected. Working the pruner ten occasions on a historical past pruned 9 occasions yields the very same end result as a clear run from scratch. This makes “run it each flip” a protected default, eliminating the necessity to observe state or purpose about earlier passes.

The one integration resolution left is how REF and DEFINE tags get connected to messages. Right here they’re literal markers contained in the message content material, which is the best mechanism for a prototype. A manufacturing system would probably connect them as structured metadata on the message object so the tags by no means leak into the uncooked textual content the mannequin reads. Go 3’s logic stays the identical both method. An upstream course of nonetheless has to find out what counts as a dependency value tagging, as a result of Go 3 can solely restore what it’s explicitly advised to trace.

What this doesn’t cowl

Dependency detection is literal, not semantic. If a reference is paraphrased and lacks an identical tag, the script will miss it.

These workloads are additionally fully artificial. I selected the three parameter units primarily based on believable manufacturing conduct, not actual telemetry. When you’ve got manufacturing logs, regenerating these three workload classes from precise utilization is the plain subsequent step. The numbers will shift relying on what your precise visitors appears to be like like.

This pipeline omits semantic compression, embeddings, and LLM-scored pruning. These are legitimate, different approaches, however avoiding them retains this implementation totally deterministic and dependency-free. LLMLingua is a primary instance of the discovered different; it makes use of a small language mannequin to attain and drop tokens, reaching a lot greater compression ratios than this script [3]. Selecting between them is a direct trade-off: you trade determinism and zero-dependency execution for tighter compression.

The token counts are additionally approximations. The script makes use of a whitespace and punctuation-boundary heuristic as an alternative of a manufacturing subword tokenizer like tiktoken. As a result of the heuristic runs persistently earlier than and after pruning, the relative discount percentages stay correct, even when absolutely the numbers don’t completely match an official tokenizer.

Lastly, there isn’t a direct latency measurement for the infrastructure causes detailed earlier.

The place I’d take this subsequent

Two extensions make sense right here somewhat than increasing the prevailing three passes.

The primary is a hybrid strategy. Preserve these three deterministic passes as a quick, protected first stage, then hand the output to an embedding-aware or LLM-scored compression device like LLMLingua. This catches semantic redundancy that literal identifier matching misses, akin to two passages saying the identical factor in several phrases.

Working the deterministic passes first preserves the protection assure. If the discovered stage misbehaves, the dialog drops again to the deterministic baseline as an alternative of failing unpredictably. Architecturally, the discovered go solely operates on the output of Go 3. A mannequin bug within the compression step may shrink a immediate too aggressively, nevertheless it can’t reintroduce a dependency failure that the deterministic passes already cleared.

The second extension is closing the hole between artificial knowledge and manufacturing actuality. The following step is regenerating these identical three workload classes from precise manufacturing traces. Maintaining the benchmark methodology similar—the identical mounted parameters, ground-truth dependency labels, and 15-configuration sweep—ensures the outcomes stay straight comparable to those printed figures whereas changing guesses with actual telemetry.

Precise utilization logs will probably shift the RAG and tool-agent metrics in both course. Actual-world visitors doesn’t robotically imply greater compression. For example, a manufacturing system with aggressive upstream deduplication may already remove the waste this artificial mannequin assumes. That might be a extremely helpful discovering in its personal proper, somewhat than a failure.

A 3rd, smaller level to notice: the REF/DEFINE conference used right here is only a placeholder. A manufacturing system ought to derive these tags robotically from structured knowledge, device name arguments, session variables, or express person settings, somewhat than counting on handbook textual content markers.

Whereas the deterministic logic in Go 3 stays similar both method, the precise worth of this pipeline relies upon fully on how cleanly you’ll be able to generate correct dependency tags upstream. That’s an architecture-specific integration drawback somewhat than one thing a general-purpose pruning library can clear up out of the field.

Manufacturing serving methods already deal with immediate bloat as a reminiscence administration drawback on the infrastructure layer, evicting and sharing KV cache pages very similar to an working system handles bodily reminiscence [4]. You possibly can consider this venture as working one layer above that, on the immediate building part, earlier than the request ever reaches the infrastructure.

The 2 layers don’t compete. A smaller, deduplicated immediate offers a greater enter to a well-managed cache somewhat than performing as a substitute for it.

The Core Takeaway

If a crucial truth is preserved, the pipeline must exhibit that retention somewhat than assume it primarily based on the absence of apparent errors. This precept guided the design of the system.

Lengthy-running conversations don’t all the time require a bigger, smarter mannequin to determine what to recollect. Typically, the extra sturdy answer is a predictable, three-pass system that may programmatically show what it didn’t lose.

Sources

[1] Levy, M., Jacoby, A., & Goldberg, Y. (2024). Identical process, extra tokens: The affect of enter size on the reasoning efficiency of huge language fashions. In Proceedings of the 62nd Annual Assembly of the Affiliation for Computational Linguistics (Quantity 1: Lengthy Papers) (pp. 15339–15353). Affiliation for Computational Linguistics. https://doi.org/10.18653/v1/2024.acl-long.818

[2] Liu, N. F., Lin, Okay., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2024). Misplaced within the center: How language fashions use lengthy contexts. Transactions of the Affiliation for Computational Linguistics, 12, 157–173. https://doi.org/10.1162/tacl_a_00638

[3] Jiang, H., Wu, Q., Lin, C.-Y., Yang, Y., & Qiu, L. (2023). LLMLingua: Compressing prompts for accelerated inference of huge language fashions. In Proceedings of the 2023 Convention on Empirical Strategies in Pure Language Processing (pp. 13358–13376). Affiliation for Computational Linguistics. https://doi.org/10.18653/v1/2023.emnlp-main.825

[4] Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Environment friendly reminiscence administration for big language mannequin serving with PagedAttention. In Proceedings of the twenty ninth Symposium on Working Techniques Ideas (SOSP ’23) (pp. 611–626). Affiliation for Computing Equipment. https://doi.org/10.1145/3600006.3613165

All code, benchmark numbers, and take a look at outcomes on this article are my very own, generated by working the included codebase straight and reproduced on two separate machines. No proprietary datasets, copyrighted textual content, or third-party code had been utilized in constructing or benchmarking this technique. All code, the corpus generator, the pruner, the benchmark harness, and all 35 assessments, is offered within the repository linked beneath.

https://github.com/Emmimal/prompt-pruning-layer