# How Needle 2 Pins Tools as KV Sinks in Memory Management

> Discover how Needle 2 pins tools as KV sinks. Learn about the boolean sink mask and bitwise OR operations that protect KV cache positions from eviction in memory management.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: internals
- Published: 2026-09-02

---

**Needle 2 pins tools as KV sinks using a boolean `sink` mask that marks tool-related KV cache positions as protected from eviction, combining this mask with the standard attention mask via bitwise OR operations in `hidden_cells` and decoder functions.**

Needle 2's memory system stores the attention key-value (KV) cache for every token generated by the model. When a tool is defined for an agent, the tool's *call* is represented as a special token sequence (e.g., `<tools>…</tools>`). During the forward pass, the model receives a **`sink` mask** — a boolean array that marks KV slots that must **never be overwritten**. This article explains exactly how this pinning mechanism works in the [cactus-compute/needle](https://github.com/cactus-compute/needle) codebase.

## The Two-Stage Pinning Mechanism

Tool pinning happens in two coordinated phases: building the mask during agent creation, then preserving pinned slots during every attention operation.

### Building the Sink Mask at Agent Initialization

When an agent is created with tools, the tool token IDs are inserted into the KV cache. A mask entry is set to `True` for those positions, and this mask is passed as the `sink` argument to the model's processing functions.

The mask construction ensures that tool definition tokens occupy protected positions before any user interaction begins. This happens in the tooling layer before the model's forward methods ever see the data.

### Preserving Sinks During Attention Computation

Every function that reads or writes KV values checks and applies the sink mask. The core merging logic lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

In the `hidden_cells` helper, the regular "recent" mask combines with the optional `sink` mask via bitwise OR:

```python
keep = recent if sink is None else (recent | sink[:, None, None, :])

```

Source: [architecture.py L541-L549](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L541-L549)

This same pattern propagates through multiple entry points. The `forward_confidence`, `encode_contrastive`, and related methods all accept a `sink` argument and pass it through:

```python
self.hidden_cells(tokens, quant=quant, window=window, sink=sink))

```

Source: [architecture.py L562-L566](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L562-L566)

## Decoder-Level Sink Application

The decoder performs the final application of sink protection before softmax-based KV updates. In [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py), the attention routine receives the same `sink` argument:

```python
keep = recent[None] if sink is None else (recent[None] | sink[:, None, :])

```

Source: [decode.py L84-L110](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py#L84-L110)

Because the `sink` mask is **OR-ed** with the regular attention mask, any KV entries belonging to pinned tools are excluded from eviction. Tool-related KV slots remain intact for the entire conversation lifespan.

## Practical Code Example

Here's how tool pinning works in practice:

```python

# 1️⃣ Create an agent with a tool definition

tools = [{"name": "send_email", "parameters": {"type": "object", "properties": {...}}}]
agent = needle.Needle(tools=tools)

# 2️⃣ Encode a prompt — the model internally builds a sink mask

prompt = "Please email the report."
logits, sink = model.encode(prompt, tools=tools)   # `sink` is a bool mask

# 3️⃣ Run a forward step — the sink mask is passed through

output = model.decode(logits, sink=sink)          # KV entries for the tool stay fixed

```

On subsequent queries, the same `sink` mask is reused. The tool's KV entries remain **pinned** and never get overwritten by newer tokens.

## Key Source Files and Functions

Understanding these files clarifies how Needle 2 implements KV sink pinning:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** — Defines the transformer core and the `hidden_cells` routine that merges the KV sink mask with the standard attention mask. [L541-L566](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L541-L566)

- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)** — Performs decoding and applies the sink mask to preserve KV slots during token generation. [L84-L110](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py#L84-L110)

- **[`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)** — Registers tool schemas and builds the JSON representation that seeds the KV sink during model initialization.

- **[`needle/playground/server.py`](https://github.com/cactus-compute/needle/blob/main/needle/playground/server.py)** — Demonstrates server-side agent recreation when tools JSON changes, ensuring fresh sink construction. [L27-L35](https://github.com/cactus-compute/needle/blob/main/needle/playground/server.py#L27-L35)

## Summary

- **KV sink pinning** protects tool definition tokens from cache eviction throughout a conversation
- The **`sink` mask** is a boolean tensor marking protected positions, built at agent initialization
- **`hidden_cells`** in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) merges sink and attention masks via `recent | sink`
- The **decoder** in [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py) applies the same merge before KV cache updates
- Tool schemas in [`tools.py`](https://github.com/cactus-compute/needle/blob/main/tools.py) provide the source data for mask construction
- Masks persist across turns, ensuring **permanent tool availability**

## Frequently Asked Questions

### What is a KV sink in Needle 2?

A KV sink is a mechanism that prevents specific key-value cache entries from being overwritten by the model's normal eviction policy. In Needle 2, tools are pinned as sinks so their token representations remain accessible for the entire conversation, enabling reliable tool calling without re-encoding definitions.

### How does the sink mask interact with attention windows?

The sink mask combines with the standard "recent" attention window mask using a bitwise OR operation (`recent | sink`). This union ensures that sink-protected positions are always included in the attention computation, regardless of whether they fall within the normal sliding window.

### Can multiple tools be pinned as KV sinks simultaneously?

Yes. The sink mask is a boolean tensor where any position can be marked `True`. When multiple tools are defined, all their respective token positions are set to `True` in the same mask, providing simultaneous protection for every tool's KV cache entries.

### When does Needle 2 rebuild the sink mask?

The sink mask rebuilds whenever the tools configuration changes. The playground server demonstrates this in [`server.py`](https://github.com/cactus-compute/needle/blob/main/server.py) lines 27-35, where agent recreation triggers fresh mask construction from updated tool schemas.