# How Tools Are Pinned as KV Sinks in Needle 2: A Technical Guide

> Learn how Needle 2 pins tools as KV sinks using a boolean mask to prevent eviction from the KV cache. Discover the technical guide for cactus-compute/needle.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-19

---

**Needle 2 prevents tool outputs from being evicted from the KV cache by accepting a boolean `sink` mask in the model's forward pass and ORing it with the sliding-window memory mask, guaranteeing tool-generated tokens remain resident across conversation steps.**

The `cactus-compute/needle` repository implements a bounded-memory language model that uses a sliding-window KV cache. To ensure built-in utilities like fetch, pydantic, and JSON-schema remain accessible throughout long conversations, Needle 2 applies a dedicated **sink mask** that marks tool-related KV positions as non-evictable. This mechanism is how tools are pinned as KV sinks in Needle 2.

## How the `@tool` Decorator Registers Functions for KV Sink Pinning

In [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), functions are decorated with `@tool`. This decorator registers the callable in the agent's tool registry and records the token IDs that will be emitted when the tool is invoked.

Before calling the model, the agent constructs a boolean array `sink` of shape `[batch, seq_len]`. Positions that correspond to tool-generated tokens are set to `True`. This pre-forward logic is exercised in [`tests/test_tools.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_tools.py).

## How the Sink Mask Flows Through the Model Architecture

### Passing the Mask Through the Public API ([`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py))

The model's public forward functions—including `hidden_cells`, `encode_contrastive`, and `forward_confidence`—accept an optional argument `sink` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The mask is forwarded down to the low-level attention kernels in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py).

### Merging the Sink Mask with the Sliding-Window Logic Inside `hidden_cells`

Inside [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `hidden_cells` method builds a "recent" mask that implements the 256-token sliding window. When a `sink` mask is supplied, the final retention mask becomes:

```python
keep = recent | sink[:, None, None, :]

```

This operation, found around line 548 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), ensures that any KV slot flagged by the sink is **never cleared**, regardless of the sliding-window budget.

## How KV Eviction Respects Pinned Tool Sinks

The KV-cache eviction logic, governed by the `KV_WINDOW_MIN` and `KV_GROUP` constants near lines 598-610 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), only discards entries where the merged mask is `False`. Because tool-related slots are always `True` after the OR operation, they stay resident and can be reused for subsequent conversation steps.

## Practical Example: Pinning Tool Tokens as KV Sinks

The following excerpt demonstrates how to construct and forward a sink mask in practice:

```python
import numpy as np

# 1. Register a tool (handled in needle/agent/tools.py)

@tool
def fetch(url: str) -> str:
    ...

# 2. Build the input and sink mask

tokens = tokenizer.encode(user_prompt + tool_call_annotation)
sink = np.zeros((batch, len(tokens)), dtype=bool)
sink[:, tool_start:tool_end] = True  # pin the tool-call tokens

# 3. Call the model with the sink mask

outputs = model.hidden_cells(tokens, sink=sink)

# 4. Inside hidden_cells (needle/model/architecture.py)

def hidden_cells(self, tokens, quant=False, window=0, sink=None):
    recent = self._recent_mask(tokens.shape[1])
    keep = recent if sink is None else (recent | sink[:, None, None, :])
    # 'keep' drives which KV slots are retained

    ...

```

## Summary

- Tools are pinned as KV sinks by registering them with the `@tool` decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) and recording their token IDs.
- A boolean `sink` mask of shape `[batch, seq_len]` marks tool-generated positions before the model forward pass.
- The `sink` mask is passed into `hidden_cells` and other forward functions in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).
- Inside the model, the mask is ORed with the sliding-window recent mask via `recent | sink[:, None, None, :]`, preventing eviction of tool KV entries.
- The eviction logic respects the merged mask, so pinned tool outputs survive the bounded-memory sliding window across multiple turns.

## Frequently Asked Questions

### What does it mean to pin a tool as a KV sink in Needle 2?

Pinning a tool as a KV sink means marking the key-value cache entries generated by a tool call as non-evictable. Needle 2 accomplishes this by passing a `sink` boolean mask into the model's forward pass, which forces the retention of those entries even when the sliding-window memory budget would normally discard older tokens.

### Which source files control the tool pinning behavior?

The pinning behavior is implemented across [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), where the `@tool` decorator registers utilities and tracks token IDs, and [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `sink` mask is merged with the sliding-window mask inside `hidden_cells`. The low-level attention kernels in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) execute the kernels that respect the merged retention mask.

### How does the sink mask interact with the 256-token sliding window?

The sink mask overrides the sliding window for specific positions. Inside [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the final retention mask is computed as `recent | sink[:, None, None, :]`, meaning any position marked `True` in the sink mask remains in the KV cache regardless of the 256-token recent-window budget.

### Can custom tools be pinned using the same mechanism?

Yes. Any function decorated with `@tool` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) can be pinned by constructing a `sink` array that sets `True` for the token positions belonging to that tool's output and passing it to the model's forward functions such as `hidden_cells`.