> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.streaming.stores.local

In-process feature store for the speculative-training stream.

The local store keeps tensors in a Python dict under the
:class:`~nemo_automodel.components.speculative.streaming.refs.SampleRef.sample_id`,
with a resident-byte counter so the queue can drive backpressure from the
same store it puts into. It is the build-and-test surface for the entire
streaming pipeline: enough to wire up a colocated producer and consumer
without a network, without a shared mount, and without GPUDirect.

Residency policy (RFC §"Open questions" Q2 answer: bytes as the hard
backstop, sample count as a soft cap). When the next `put` would exceed
either cap, :meth:`put` raises :class:`MemoryError` -- the producer is then
expected to retry after the consumer drains below the low watermark.

## Module Contents

### Classes

| Name                                                                                                   | Description                                      |
| ------------------------------------------------------------------------------------------------------ | ------------------------------------------------ |
| [`LocalFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-local-LocalFeatureStore) | In-process :class:`FeatureStore` implementation. |

### Data

[`__all__`](#nemo_automodel-components-speculative-streaming-stores-local-__all__)

[`logger`](#nemo_automodel-components-speculative-streaming-stores-local-logger)

### API

```python
class nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore(
    max_samples: int | None = 64,
    max_bytes: int | None = 256 * 1024 * 1024,
    high_watermark_bytes: int | None = 192 * 1024 * 1024,
    low_watermark_bytes: int | None = 64 * 1024 * 1024
)
```

**Bases:** [FeatureStore](/nemo-automodel/nemo_automodel/components/speculative/streaming/store#nemo_automodel-components-speculative-streaming-store-FeatureStore)

In-process :class:`FeatureStore` implementation.

Thread safety: every public method holds a single :class:`threading.Lock`,
so concurrent puts and gets from the same Python process are safe. Async
/ cross-process safety is the queue's responsibility and is out of scope
for PR 1.

**Parameters:**

**`max_samples`** `int | None` — default: 64

Hard cap on simultaneously-stored samples. `None` means
unbounded sample count (still bounded by `max_bytes`).

---

**`max_bytes`** `int | None` — default: 256 \* 1024 \* 1024

Hard cap on resident bytes. `None` means unbounded
(still bounded by `max_samples`). At least one of
`max_samples` / `max_bytes` must be set, otherwise a
misconfigured store silently behaves as unbounded.

---

**`high_watermark_bytes`** `int | None` — default: 192 \* 1024 \* 1024

Threshold above which :attr:`StoreHealth.high_watermark_hit`
is `True`. The producer pauses here.

---

**`low_watermark_bytes`** `int | None` — default: 64 \* 1024 \* 1024

Threshold below which :attr:`StoreHealth.low_watermark_hit`
is `True`. The producer resumes here. Must be
strictly less than `high_watermark_bytes`; a hysteresis band
of zero flaps the producer on every step.

---

**`_handle_refs`** `dict[int, str] = {}`

---

**`_high_watermark`**

---

**`_lock`** `= threading.Lock()`

---

**`_low_watermark`**

---

**`_resident_bytes`** `= 0`

---

**`_storage`** `dict[str, dict[str, Tensor]] = {}`

---

**`_uri_id`** `= uuid.uuid4().hex`

---

**`store_uri`** `str`

---

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore._make_ref(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    schema_version: int,
    target_model_version: str,
    draft_weight_version: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm,
    num_tokens: int
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore._tensor_bytes(
    tensor: torch.Tensor
) -> int
```

staticmethod

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.close() -> None
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.gc() -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.get(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    device: torch.device | str | None = None
) -> tuple[dict[str, torch.Tensor], nemo_automodel.components.speculative.streaming.store.StoreHandle]
```

Materialize `ref`'s features on `device` and hand back a :class:`StoreHandle`.

**Parameters:**

**`ref`** `SampleRef`

The reference returned by :meth:`put` (typically via a
queue lease). `ref.store_uri` MUST equal this store's
:attr:`store_uri`; a mismatch raises `KeyError` so a
consumer cannot accidentally materialize a foreign ref.

---

**`device`** `torch.device | str | None` — default: None

Optional target device. `None` returns each feature
on the device it was put on; a non-`None` value
materializes every feature on that device via
`Tensor.to(device)` (a no-op when already in place).

---

**Returns:** `dict[str, torch.Tensor]`

A `(tensors, handle)` pair. `tensors` is a

**Raises:**

* `KeyError`: when `ref.store_uri` does not match this store,
  or when `ref.sample_id` is no longer present (released
  or never put).
* `RuntimeError`: when the stored tensor's shape or dtype
  differs from what the ref claims, or the store has been
  closed.

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm = FeatureAlgorithm.EAGLE3,
    schema_version: int = 1,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    num_tokens: int = 0
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

Store `tensors` under `sample_id` and return a tensor-free :class:`SampleRef`.

**Parameters:**

**`sample_id`** `str`

Stable identifier within `run_id`. Must be unique
in this store at put time; duplicates raise `ValueError`.

---

**`tensors`** `Mapping[str, torch.Tensor]`

Feature-name to tensor mapping. The store detaches,
clones, and makes each tensor contiguous before stashing it,
so the producer may keep mutating its source tensors after
the put returns without disturbing what a later
:meth:`get` hands out. The shape and dtype of each tensor
are captured into the returned :class:`SampleRef`'s
`feature_specs`; the consumer uses those specs to
preallocate the receive buffer at :meth:`get` time, so
changing `tensors[name].shape` or `dtype` between put
and get without updating the ref will surface as a
`RuntimeError` on materialization.

---

**`run_id`** `str`

Same value on every ref of one run; surfaces on the
:class:`SampleRef.run_id` so producers and consumers can
verify they are talking about the same run.

---

**`algorithm`** `FeatureAlgorithm` — default: FeatureAlgorithm.EAGLE3

Which draft family produced this sample; gates the
:class:`SampleRef` required-features check.

---

**`schema_version`** `int` — default: 1

Bumped whenever the producer's feature set or
layout for `algorithm` changes incompatibly.

---

**`target_model_version`** `str` — default: '0'

Monotonically increasing identifier of
the target-model weights.

---

**`draft_weight_version`** `str` — default: '0'

Same idea for the draft model's weights.

---

**`num_tokens`** `int` — default: 0

Sum of attended tokens; used by the consumer for
empty / short loss-mask neutralization.

---

**Returns:** `SampleRef`

A tensor-free :class:`SampleRef` carrying the per-feature

**Raises:**

* `MemoryError`: if the put would exceed `max_samples` or
  `max_bytes`. The producer is expected to retry after the
  store drains below the low watermark (see
  :meth:`health`).
* `RuntimeError`: if the store has been closed.
* `ValueError`: on bad input (empty sample id, empty tensors map,
  duplicate sample id).

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.release(
    handle: nemo_automodel.components.speculative.streaming.store.StoreHandle
) -> None
```

```python
nemo_automodel.components.speculative.streaming.stores.local.__all__ = ['LocalFeatureStore']
```

```python
nemo_automodel.components.speculative.streaming.stores.local.logger = logging.getLogger(__name__)
```