> 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.refs

Tensor-free control-plane contracts for the speculative-training stream.

Per the train-inference disaggregation RFC (issue #3062), every produced sample
is split into a tensor-free reference (this module) and the actual supervision
tensors, which live in a pluggable :mod:`store`. References hop between the
target-side producer and the draft-side consumer over queues, HTTP, and
checkpoint metadata -- places where serializing a :class:`torch.Tensor` would be
either expensive or wrong.

Three guarantees back this module:

1. :func:`assert_no_tensors` enforces "no tensors" recursively on every public
   object exposed from this package, so a producer that slipped a tensor into a
   queue or ref trips a validation error before it lands on a wire.
2. :class:`SampleRef` is a frozen dataclass. Once placed on a queue, a ref
   cannot be mutated under the holder, so the consumer is guaranteed to see the
   same feature keys the producer promised.
3. :class:`FeatureSpec` carries `dtype` + `shape` (and nothing else), the
   minimum metadata a consumer needs to preallocate the receive buffer *before*
   materializing the tensor -- mirroring how
   :func:`nemo_automodel.components.speculative.eagle.remote.protocol.encode_nccl_metadata`
   already ships dtype + shape ahead of an NCCL recv so the client can allocate.

## Module Contents

### Classes

| Name                                                                                         | Description                                                           |
| -------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`FeatureAlgorithm`](#nemo_automodel-components-speculative-streaming-refs-FeatureAlgorithm) | Which speculative-decoding draft family produced this sample.         |
| [`FeatureSpec`](#nemo_automodel-components-speculative-streaming-refs-FeatureSpec)           | Shape + dtype metadata for one named feature in a :class:`SampleRef`. |
| [`SampleRef`](#nemo_automodel-components-speculative-streaming-refs-SampleRef)               | Tensor-free reference to one produced sample in a feature store.      |

### Functions

| Name                                                                                                                 | Description                                                                   |
| -------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_algorithm_required_features`](#nemo_automodel-components-speculative-streaming-refs-_algorithm_required_features) | Required feature-name set per algorithm (RFC "Feature schema per algorithm"). |
| [`_is_numpy_array`](#nemo_automodel-components-speculative-streaming-refs-_is_numpy_array)                           | -                                                                             |
| [`_is_torch_tensor`](#nemo_automodel-components-speculative-streaming-refs-_is_torch_tensor)                         | -                                                                             |
| [`_reject_tensor`](#nemo_automodel-components-speculative-streaming-refs-_reject_tensor)                             | -                                                                             |
| [`assert_no_tensors`](#nemo_automodel-components-speculative-streaming-refs-assert_no_tensors)                       | Recursively validate that `obj` carries no tensors.                           |

### Data

[`_PRIMITIVE_TYPES`](#nemo_automodel-components-speculative-streaming-refs-_PRIMITIVE_TYPES)

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

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

### API

```python
class nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm
```

**Bases:** `enum.Enum`

Which speculative-decoding draft family produced this sample.

The on-wire string value (`eagle3` / `dflash` / `dspark`) is the
stable schema key the consumer matches against. New algorithms land by
adding a member here plus a schema row in the RFC's "Feature schema per
algorithm" table; both moves are part of the same change.

**`DFLASH`** `= 'dflash'`

---

**`DSPARK`** `= 'dspark'`

---

**`EAGLE3`** `= 'eagle3'`

---

```python
class nemo_automodel.components.speculative.streaming.refs.FeatureSpec(
    shape: tuple[int, ...],
    dtype: torch.dtype
)
```

Dataclass

Shape + dtype metadata for one named feature in a :class:`SampleRef`.

Mirrors the `{dtype_code, shape}` pair the EAGLE-3 remote protocol
encodes in :func:`nemo_automodel.components.speculative.eagle.remote.protocol.encode_nccl_metadata`
so the client can preallocate the NCCL receive buffer. Consumers here use
it to preallocate the receive tensor before calling
:meth:`FeatureStore.get`.

**`dtype`** `dtype`

---

**`shape`** `tuple[int, ...]`

---

```python
nemo_automodel.components.speculative.streaming.refs.FeatureSpec.__post_init__() -> None
```

```python
class nemo_automodel.components.speculative.streaming.refs.SampleRef(
    sample_id: str,
    run_id: str,
    store_uri: str,
    feature_keys: dict[str, str],
    feature_specs: dict[str, nemo_automodel.components.speculative.streaming.refs.FeatureSpec],
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm,
    schema_version: int,
    num_tokens: int,
    estimated_bytes: int,
    target_model_version: str,
    draft_weight_version: str
)
```

Dataclass

Tensor-free reference to one produced sample in a feature store.

**`algorithm`** `FeatureAlgorithm`

---

**`draft_weight_version`** `str`

---

**`estimated_bytes`** `int`

---

**`feature_keys`** `dict[str, str]`

---

**`feature_specs`** `dict[str, FeatureSpec]`

---

**`num_tokens`** `int`

---

**`run_id`** `str`

---

**`sample_id`** `str`

---

**`schema_version`** `int`

---

**`store_uri`** `str`

---

**`target_model_version`** `str`

---

```python
nemo_automodel.components.speculative.streaming.refs.SampleRef.__post_init__() -> None
```

```python
nemo_automodel.components.speculative.streaming.refs.SampleRef.feature_names() -> tuple[str, ...]
```

Return the feature names in insertion order.

The order is significant because some algorithms size the draft's
`fc` projection by the *order* of the aux hidden states, not just
their count; the producer chooses the order once, the ref freezes it,
and the consumer preserves it through :meth:`FeatureStore.get`.

```python
nemo_automodel.components.speculative.streaming.refs._algorithm_required_features(
    algo: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm
) -> frozenset[str]
```

Required feature-name set per algorithm (RFC "Feature schema per algorithm").

Mirrors the "Required feature keys" column of the RFC's schema table.
PR 2 widens this into a typed registry that also validates per-key
dtypes / shapes; PR 1 only checks that the producer named the right set.

```python
nemo_automodel.components.speculative.streaming.refs._is_numpy_array(
    obj: typing.Any
) -> bool
```

```python
nemo_automodel.components.speculative.streaming.refs._is_torch_tensor(
    obj: typing.Any
) -> bool
```

```python
nemo_automodel.components.speculative.streaming.refs._reject_tensor(
    obj: typing.Any,
    path: str
) -> typing.NoReturn
```

```python
nemo_automodel.components.speculative.streaming.refs.assert_no_tensors(
    obj: typing.Any,
    path: str = 'ref'
) -> None
```

Recursively validate that `obj` carries no tensors.

Walks dataclasses (any nested dataclass included), `dict` with `str`
keys, and `list` / `tuple`. Anything else must be a primitive
(`str`/`int`/`float`/`bool`/`None`/`bytes`) -- a tensor, numpy
array, or duck-typed tensor-like (`has data_ptr + is_cuda + numel`) at
any depth raises :class:`ValueError`.

The check is structural, not nominal: a third-party tensor type that
quacks like one is rejected. PR 2's registry / schema validators layer on
top of this primitive guard, not instead of it.

```python
nemo_automodel.components.speculative.streaming.refs._PRIMITIVE_TYPES = (str, int, float, bool, type(None), bytes)
```

```python
nemo_automodel.components.speculative.streaming.refs.__all__ = ['FeatureAlgorithm', 'FeatureSpec', 'SampleRef', 'assert_no_tensors']
```

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