> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

Pluggable feature-store backends for the streaming data plane.

`LocalFeatureStore` is the in-process reference backend.
`SharedDirFeatureStore` adds cross-process rendezvous on a shared
POSIX mount. A future `NcclFeatureStore` can plug in behind the same
`FeatureStore`
contract.

## Submodules

* **[`nemo_automodel.components.speculative.streaming.stores.local`](/nemo-automodel/nemo_automodel/components/speculative/streaming/stores/local)**
* **[`nemo_automodel.components.speculative.streaming.stores.shared_dir`](/nemo-automodel/nemo_automodel/components/speculative/streaming/stores/shared_dir)**

## Package Contents

### Classes

| Name                                                                                                                | Description                                                              |
| ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`LocalFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-local-LocalFeatureStore)              | In-process `FeatureStore` implementation.                                |
| [`SharedDirFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-SharedDirFeatureStore) | `FeatureStore` backed by one `<sample_id>.safetensors` per sample. |

### 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 `FeatureStore` implementation.

Thread safety: every public method holds a single `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 `StoreHealth.high_watermark_hit`
is `True`. The producer pauses here.

---

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

Threshold below which `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 `StoreHandle`.

**Parameters:**

**`ref`** `SampleRef`

The reference returned by `put` (typically via a
queue lease). `ref.store_uri` MUST equal this store's
`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 `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
`get` hands out. The shape and dtype of each tensor
are captured into the returned `SampleRef`'s
`feature_specs`; the consumer uses those specs to
preallocate the receive buffer at `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
`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
`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 `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
  `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
class nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore(
    directory: str,
    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)

`FeatureStore` backed by one `<sample_id>.safetensors` per sample.

Thread safety: every public method holds a single
`threading.Lock`. Cross-process / cross-node is supported
as long as distinct processes use distinct `sample_id` values;
the lock does not extend across processes.

**Parameters:**

**`directory`** `str`

Filesystem path used as the rendezvous. Created if it
does not exist. Concurrent producers and consumers in
separate processes / ranks coordinate via unique
`sample_id` values -- collision is the caller's problem.

---

**`max_samples, max_bytes, high_watermark_bytes, low_watermark_bytes`**

Same residency contract as `LocalFeatureStore`; the
queue's HWM/LWM hysteresis reads them off
`health`.

---

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

---

**`_high_watermark`**

---

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

---

**`_low_watermark`**

---

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

---

**`store_uri`** `str`

---

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._atomic_write(
    path: str,
    tensors: typing.Mapping[str, torch.Tensor]
) -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._path_for(
    sample_id: str
) -> str
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.close() -> None
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.gc() -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.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]
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.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
```

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