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:

: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.
: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.
: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`	Which speculative-decoding draft family produced this sample.
`FeatureSpec`	Shape + dtype metadata for one named feature in a :class:`SampleRef`.
`SampleRef`	Tensor-free reference to one produced sample in a feature store.

Functions

Name	Description
`_algorithm_required_features`	Required feature-name set per algorithm (RFC “Feature schema per algorithm”).
`_is_numpy_array`	-
`_is_torch_tensor`	-
`_reject_tensor`	-
`assert_no_tensors`	Recursively validate that `obj` carries no tensors.

Data

_PRIMITIVE_TYPES

__all__

logger

API

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'

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

shape

tuple[int, ...]

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

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

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

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.

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.

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

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

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

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.

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

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

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