> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/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 Tensor-free control-plane contracts for the speculative-training stream. The streaming pipeline (per the EAGLE-3 / DFlash / DSpark train-inference disaggregation RFC, issue #3062 PR 1) splits every transferred sample into a `refs` (control-plane, no tensors) and a `store` (data-plane, holds the supervision tensors). The queue carries refs only; tensors live in the store and are referenced by `nemo_automodel.components.speculative.streaming.refs.SampleRef.feature_keys`. This package owns: * `SampleRef` and `FeatureSpec` \-- the frozen, tensor-free reference carried on every control-plane hop. * `assert_no_tensors` \-- the guard that enforces the no-tensor invariant on the control plane. * `FeatureStore` \-- the pluggable data-plane transport (local dict, shared POSIX mount, NCCL, ...); `LocalFeatureStore` \-- the in-process implementation land-tested by PR 1. * `SampleRefQueue` \-- the metadata-only lease/ack/fail queue between producers and consumers with visibility-timeout reclaim and watermark-based backpressure. * `FeatureProducer` and `FeatureDataLoader` \-- the EAGLE-3 produce and consume sides that turn a target forward into refs and refs back into `Eagle3TargetBatch` instances. * `AsyncFeaturePipeline` \-- runs the producer in a background thread so target forward and draft backward overlap; `SharedDirFeatureStore` adds a POSIX shared-mount data plane for cross-process producer/consumer. ## Subpackages * **[`nemo_automodel.components.speculative.streaming.stores`](/nemo-automodel/nemo_automodel/components/speculative/streaming/stores)** ## Submodules * **[`nemo_automodel.components.speculative.streaming.async_pipeline`](/nemo-automodel/nemo_automodel/components/speculative/streaming/async_pipeline)** * **[`nemo_automodel.components.speculative.streaming.eagle3`](/nemo-automodel/nemo_automodel/components/speculative/streaming/eagle3)** * **[`nemo_automodel.components.speculative.streaming.loader`](/nemo-automodel/nemo_automodel/components/speculative/streaming/loader)** * **[`nemo_automodel.components.speculative.streaming.producer`](/nemo-automodel/nemo_automodel/components/speculative/streaming/producer)** * **[`nemo_automodel.components.speculative.streaming.queue`](/nemo-automodel/nemo_automodel/components/speculative/streaming/queue)** * **[`nemo_automodel.components.speculative.streaming.refs`](/nemo-automodel/nemo_automodel/components/speculative/streaming/refs)** * **[`nemo_automodel.components.speculative.streaming.store`](/nemo-automodel/nemo_automodel/components/speculative/streaming/store)** ## Package Contents ### Classes | Name | Description | | ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- | | [`AsyncFeaturePipeline`](#nemo_automodel-components-speculative-streaming-async_pipeline-AsyncFeaturePipeline) | Run a `FeatureProducer` in a background thread, draining a prompt source. | | [`FeatureAlgorithm`](#nemo_automodel-components-speculative-streaming-refs-FeatureAlgorithm) | Which speculative-decoding draft family produced this sample. | | [`FeatureDataLoader`](#nemo_automodel-components-speculative-streaming-loader-FeatureDataLoader) | Iterator over `Eagle3TargetBatch` materialized from a streaming queue. | | [`FeatureProducer`](#nemo_automodel-components-speculative-streaming-producer-FeatureProducer) | Run a target backend once and put the supervision into a `FeatureStore`. | | [`FeatureSpec`](#nemo_automodel-components-speculative-streaming-refs-FeatureSpec) | Shape + dtype metadata for one named feature in a `SampleRef`. | | [`FeatureStore`](#nemo_automodel-components-speculative-streaming-store-FeatureStore) | Abstract transport every feature-store backend implements. | | [`Lease`](#nemo_automodel-components-speculative-streaming-queue-Lease) | Handle to a leased `SampleRef`. | | [`LocalFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-local-LocalFeatureStore) | In-process `FeatureStore` implementation. | | [`PromptSource`](#nemo_automodel-components-speculative-streaming-async_pipeline-PromptSource) | Zero-argument callable that supplies the next prompt batch. | | [`SampleRef`](#nemo_automodel-components-speculative-streaming-refs-SampleRef) | Tensor-free reference to one produced sample in a feature store. | | [`SampleRefQueue`](#nemo_automodel-components-speculative-streaming-queue-SampleRefQueue) | Lease / ack / fail queue over `SampleRef`. | | [`SharedDirFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-SharedDirFeatureStore) | `FeatureStore` backed by one `<sample_id>.safetensors` per sample. | | [`StoreHandle`](#nemo_automodel-components-speculative-streaming-store-StoreHandle) | Opaque token the consumer must return to `FeatureStore.release`. | | [`StoreHealth`](#nemo_automodel-components-speculative-streaming-store-StoreHealth) | Integers-only snapshot of store residency for backpressure decisions. | | [`VisibilityTimeout`](#nemo_automodel-components-speculative-streaming-queue-VisibilityTimeout) | How long an unacked `Lease` is allowed to live before reclaim. | ### Functions | Name | Description | | ---------------------------------------------------------------------------------------------- | --------------------------------------------------- | | [`assert_no_tensors`](#nemo_automodel-components-speculative-streaming-refs-assert_no_tensors) | Recursively validate that `obj` carries no tensors. | ### API ```python class nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline( producer: nemo_automodel.components.speculative.streaming.producer.FeatureProducer, queue: nemo_automodel.components.speculative.streaming.queue.SampleRefQueue, prompt_source: nemo_automodel.components.speculative.streaming.async_pipeline.PromptSource | typing.Iterator[nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | torch.Tensor], poll_interval: float = 0.1, stop_on_exhausted: bool = True ) ``` Run a `FeatureProducer` in a background thread, draining a prompt source. **Lifecycle** Construct, then `start` (or use the context manager). The background thread is daemon; `stop` joins it within `join_timeout` seconds. Outstanding leases are the trainer's responsibility -- the pipeline does not drop them on close. **Parameters:** **`producer`** `FeatureProducer` The `FeatureProducer` to invoke on each prompt. It carries the wrapped target backend, the store, and the per-call metadata. --- **`queue`** `SampleRefQueue` The `SampleRefQueue` to push the resulting `SampleRef` onto. The queue's HWM/LWM hysteresis paces the producer against the consumer. --- **`prompt_source`** `PromptSource | Iterator[_PromptBatch | _PackedPromptBatch | torch.Tensor]` A `PromptSource` -- either a zero-arg callable or an `Iterator`. The callable is invoked from the background thread; it must be thread-safe (e.g. a `torch.utils.data.DataLoader` iterator is). --- **`poll_interval`** `float` — default: 0.1 Seconds between `prompt_source` invocations after exhaustion when `stop_on_exhausted` is `False`. Defaults to 100ms -- cheap, lets the producer resume quickly if more data lands. --- **`stop_on_exhausted`** `bool` — default: True When `True` (default), the background thread exits as soon as `prompt_source` returns `None` and `close` drains. When `False`, the thread keeps polling so a streaming dataset can refill. --- **`_error`** `BaseException | None = None` --- **`_iterator`** `Iterator[_PromptBatch | _PackedPromptBatch | Tensor] | None = prompt_source` --- **`_prompt_source`** `Callable[[], _PromptBatch | _PackedPromptBatch | None] = self._pull_from_iterator` --- **`_stop_event`** `= threading.Event()` --- **`_thread`** `Thread | None = None` --- ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.__enter__() -> 'AsyncFeaturePipeline' ``` ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.__exit__( exc_type, exc, tb ) -> None ``` ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline._pull_from_iterator() -> nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | None ``` ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline._run() -> None ``` ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.join( timeout: float | None = None ) -> None ``` ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.start() -> None ``` Spawn the background producer thread. Idempotent while the thread is alive: a second call is a no-op. The thread is named `streaming-async-producer` so test failures and runtime traces show which thread is hung. The pipeline is single-use. When the producer thread exits -- prompt source exhausted, `stop`, or an error -- the bound queue is closed in `_run`'s `finally` so the consumer drains and stops. A closed queue cannot reopen, so restarting is not supported; construct a fresh pipeline (and queue) to stream again. Restart is rejected here rather than failing later with an opaque "queue is closed" from a blocking put. ```python nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.stop( join_timeout: float | None = 10.0 ) -> None ``` Signal the background thread to exit and join it. Outstanding leases are the trainer's; this method does not ack them. Idempotent. If the background thread raises, the exception is re-raised here so the trainer sees the failure. ```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.loader.FeatureDataLoader( queue: nemo_automodel.components.speculative.streaming.queue.SampleRefQueue, store: nemo_automodel.components.speculative.streaming.store.FeatureStore, algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm | None = None, acquire_poll_interval: float = 0.05 ) ``` Iterator over `Eagle3TargetBatch` materialized from a streaming queue. **Lifecycle** Each iterator pull yields an `Eagle3TargetBatch` whose tensors come from a fresh `store.get()`. The previous batch's lease is ack'd and its store handle released on the NEXT pull -- so the trainer can hold one batch across one forward pass without it being freed mid-forward. A consumer that wants eager reclamation (e.g. to free memory before pulling the next batch) calls `consume_now` after computing its loss. Iteration ends on `close` or when the queue drains. **Parameters:** **`queue`** `SampleRefQueue` The metadata-only queue the producer puts `SampleRef` onto. `queue.close()` at any time cuts the iterator short. --- **`store`** `FeatureStore` The `FeatureStore` each lease will be materialized through. Must match `ref.store_uri` for the leased refs. --- **`algorithm`** `FeatureAlgorithm | None` — default: None `FeatureAlgorithm` the loader runs the per-algorithm schema check for; defaults to EAGLE-3. --- **`_pending_handle`** `StoreHandle | None = None` --- **`_pending_lease`** `Lease | None = None` --- ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__enter__() ``` ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__exit__( exc_type, exc, tb ) ``` ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__iter__() ``` ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__next__() -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch ``` ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader._release_pending() -> None ``` ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.close() -> None ``` Ack the most recent lease / release its store handle and shut the queue. After `close` the iterator raises `StopIteration` on the next pull. Idempotent so a trainer can use `with` safely. ```python nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.consume_now() -> None ``` Release the most recently yielded batch eagerly. Useful for trainer hooks that want to free memory before pulling the next batch (e.g. immediately after `backward()`). Idempotent. ```python class nemo_automodel.components.speculative.streaming.producer.FeatureProducer( target_backend: nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend, store: nemo_automodel.components.speculative.streaming.store.FeatureStore, run_id: str, algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm | None = None, target_model_version: str = '0', draft_weight_version: str = '0', sample_id_factory: typing.Callable[[int], str] | None = None ) ``` Run a target backend once and put the supervision into a `FeatureStore`. Thread safety: `produce` is not reentrant with itself on the same producer instance; one thread should drive `produce` at a time. Multiple producers over the same store are fine; the store's lock serializes the puts. **Parameters:** **`target_backend`** `Eagle3TargetBackend` an object exposing `generate_batch(input_ids, attention_mask, loss_mask, ...) -> Eagle3TargetBatch` plus the `get_input_embeddings` and `set_vocab_mapping` accessors from `Eagle3TargetBackend`. --- **`store`** `FeatureStore` The `FeatureStore` every `produce` call writes into. --- **`run_id`** `str` Stable across the whole training run; mirrored onto every `SampleRef`. --- **`algorithm`** `FeatureAlgorithm | None` — default: None Forced `FeatureAlgorithm`; defaults to picking by the backend's runtime type. --- **`target_model_version`** `str` — default: '0' Monotonically increasing identifier of the target-model weights; surfaced on each `SampleRef`. --- **`draft_weight_version`** `str` — default: '0' Same idea for the draft model's weights. --- **`sample_id_factory`** `Callable[[int], str] | None` — default: None Callable that produces the per-call sample id; defaults to `"sample-{run_id}-{n}"`. Wrap with a hash of the input batch if you want stable ids across runs. --- **`_algorithm`** --- **`_sample_id_factory`** `= sample_id_factory or self._default_sample_id` --- **`_sample_seq`** `= 0` --- ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer._default_sample_id( n: int ) -> str ``` ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer.close() -> None ``` Release any backend resources; the producer itself holds none. ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer.discard( ref: nemo_automodel.components.speculative.streaming.refs.SampleRef ) -> None ``` Evict a produced sample whose ref will not be enqueued. Used when a produced ref is dropped before it reaches the queue: a spec-mismatch abort in `produce`, or the async pipeline shutting down while a blocking put is in flight. `get` + `release` drops the last handle so the store frees the sample; without it the sample's bytes (or its shared-dir file) stay resident with no consumer able to release them. **Parameters:** **`ref`** `SampleRef` The reference whose backing sample should be evicted. --- ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer.get_input_embeddings() -> torch.nn.Module ``` Expose the wrapped backend's input-embedding module. Mirrors `Eagle3TargetBackend.get_input_embeddings` so the draft can seed its input embeddings from the target before the first streamed batch arrives. ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer.produce( input_ids: torch.Tensor, attention_mask: torch.Tensor, loss_mask: torch.Tensor, position_ids: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None, doc_remaining: torch.Tensor | None = None ) -> nemo_automodel.components.speculative.streaming.refs.SampleRef ``` Run the wrapped backend once and stash its supervision into the store. **Parameters:** **`input_ids`** `torch.Tensor` Tensor of shape `[batch, sequence]`, `torch.long`. --- **`attention_mask`** `torch.Tensor` Tensor of shape `[batch, sequence]`, `torch.long`. --- **`loss_mask`** `torch.Tensor` Tensor of shape `[batch, sequence]`, `torch.long`. --- **`position_ids`** `torch.Tensor | None` — default: None Per-document position ids when packing is on; `None` for non-packed. Carried through unchanged to the `Eagle3TargetBatch` so the trainer's loss path stays bit-identical to the colocated path. --- **`seq_lens`** `torch.Tensor | None` — default: None Per-document sequence lengths when packing is on; `None` for non-packed. --- **`doc_remaining`** `torch.Tensor | None` — default: None `[batch, sequence]` long tensor gating cross-document TTT supervision when packing is on; `None` otherwise. --- **Returns:** `SampleRef` class:`SampleRef` carrying a tensor-free description of **Raises:** * `MemoryError`: forwarded from the store when the put would exceed `max_samples` or `max_bytes`. Producers can wait for the store to drain through the queue's backpressure path (`put_blocks_until_below`). ```python nemo_automodel.components.speculative.streaming.producer.FeatureProducer.set_vocab_mapping( selected_token_ids: torch.Tensor, selected_token_mask: torch.Tensor ) -> None ``` Thread the draft-vocab mapping through to the wrapped backend. The colocated backend ignores the mapping and projects the full-vocab `logits` trainer-side; a future draft-vocab encoding producer supplies the mapping here so the target can precompute `target_probs` + `position_mask` and the wire never carries a full-vocab tensor. ```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 `SampleRef`. Mirrors the `{dtype_code, shape}` pair the EAGLE-3 remote protocol encodes in `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 `FeatureStore.get`. **`dtype`** `dtype` The `torch.dtype` the producer will store under `SampleRef.feature_keys[name]`. Validation in `assert_no_tensors` rejects stray tensors in the ref but a `torch.dtype` object on this field is allowed. --- **`shape`** `tuple[int, ...]` Logical shape with the same axis order the producer used. `shape[0]` is conventionally batch and `shape[-1]` the feature dimension; algorithms with packed layouts (THD, BHSD) document their conventions in the algorithm row of the RFC's "Feature schema per algorithm" table. --- ```python nemo_automodel.components.speculative.streaming.refs.FeatureSpec.__post_init__() -> None ``` ```python class nemo_automodel.components.speculative.streaming.store.FeatureStore() ``` Abstract Abstract transport every feature-store backend implements. Implementations must be safe to call from one thread per operation; cross-thread concurrency is the queue's responsibility and outside this contract. The `SampleRefQueue` uses `health` for backpressure; everything else is the produce/consume pair plus `release` / `gc` for lifecycle. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.close() -> None ``` abstract Release every resource owned by the store (file handles, NCCL groups, ...). After `close`, all subsequent calls raise `RuntimeError`. Idempotent with respect to an already-closed store. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.gc() -> int ``` abstract Sweep stale entries (e.g. failed releases from a crashed consumer). Returns the number of entries reclaimed. Called opportunistically by the `SampleRefQueue` between leases and unconditionally by `close`. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.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] ``` abstract Materialize `ref`'s tensors on `device` and hand back a `StoreHandle`. The returned tensors are detached copies (`clone` for `cuda` tensors, plain `detach` for CPU views) so prefetch cannot observe aliasing through `release`. Returns one tensor per key in `SampleRef.feature_keys`, in the same insertion order, so the consumer can wire them straight into the per-algorithm batch. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth ``` abstract Return a ints-only `StoreHealth` snapshot for backpressure. Must not block on tensor I/O (no `.cpu()`, no `.to()`); in-memory counters are sufficient. Backends that do background I/O MUST serve `health` from cached counters, not from the in-flight I/O thread. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.put( sample_id: str, tensors: typing.Mapping[str, torch.Tensor], run_id: str, algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm, schema_version: int, target_model_version: str, draft_weight_version: str, num_tokens: int ) -> nemo_automodel.components.speculative.streaming.refs.SampleRef ``` abstract Stash `tensors` for `sample_id` and return a tensor-free `SampleRef`. The producer-side metadata (`run_id`, `algorithm`, `schema_version`, `target_model_version`, `draft_weight_version`, `num_tokens`) is the data the `SampleRef` carries on the control plane, so the store must accept it here even though it does not inspect the values beyond building the ref. Consumers see the ref unchanged. Implementations must validate that the tensors match what they declare (`dtype`, `shape` per feature, `numel * element_size` summed across features == `ref.estimated_bytes`) and reject the put with a specific exception when they do not, before any partial write is observable. ```python nemo_automodel.components.speculative.streaming.store.FeatureStore.release( handle: nemo_automodel.components.speculative.streaming.store.StoreHandle ) -> None ``` abstract Free the resources backing `handle`; idempotent. After this returns, the sample is no longer present in the store and a second `get` on the same `SampleRef` raises `KeyError`. `gc` retries releases that a previous call rejected (e.g. transient I/O), so the queue can rely on `gc + release` being idempotent + retriable. ```python class nemo_automodel.components.speculative.streaming.queue.Lease( ref: nemo_automodel.components.speculative.streaming.refs.SampleRef, deadline: float, visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout, redelivery_count: int = 0, lease_id: int = _next_lease_id() ) ``` Dataclass Handle to a leased `SampleRef`. Each `SampleRefQueue.acquire` call mints a fresh `Lease` with a unique `lease_id`. The queue's `ack` and `fail` verify the lease identity before mutating internal state, so a late ACK for a stale (reclaimed) lease cannot pop a newer active lease for the same `sample_id`. **`deadline`** `float` Monotonic-clock timestamp at which this lease is considered orphaned. Used by `SampleRefQueue.reclaim_expired` to redeliver the ref. --- **`lease_id`** `int = field(default_factory=_next_lease_id)` Per-acquire unique identifier. The queue uses it as the key in `_outstanding` and verifies it before any ack/fail mutation. --- **`redelivery_count`** `int = 0` Number of times this ref has been leased and re-leased (used for retry telemetry). Starts at 0. --- **`ref`** `SampleRef` The leased reference -- the only sanctioned way to materialize its tensors is `FeatureStore.get`, which returns a `StoreHandle` the consumer must hand to `FeatureStore.release` once it is done with them. --- **`visibility_timeout`** `VisibilityTimeout` The `VisibilityTimeout` that produced this lease, kept here so the consumer can introspect it. --- ```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.async_pipeline.PromptSource() ``` Protocol Zero-argument callable that supplies the next prompt batch. ```python nemo_automodel.components.speculative.streaming.async_pipeline.PromptSource.__call__() -> nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | None ``` Return the next prompt batch or `None` when the source is exhausted. **Returns:** `_PromptBatch | _PackedPromptBatch | None` A three-tuple `(input_ids, attention_mask, loss_mask)` where ```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` Which draft family produced this sample; the consumer matches it against a registered `FeatureSpec` registry to validate the ref before materializing. --- **`draft_weight_version`** `str` Same idea for the draft model's weights, so a consumer can refuse to train against a ref produced before its own weight snapshot. Starts at `"0"`; PR 4 wires the resync. --- **`estimated_bytes`** `int` Sum of every feature's `numel() * dtype_bytes`. The `SampleRefQueue` reads this against `StoreHealth.resident_bytes` for watermark hysteresis. --- **`feature_keys`** `dict[str, str]` Named features this sample contributes, mapped to per-store keys (e.g. filenames for `SharedDirFeatureStore` in PR 3, dict keys for `LocalFeatureStore`). --- **`feature_specs`** `dict[str, FeatureSpec]` Per-feature `FeatureSpec` so the consumer can allocate the receive buffer before calling `FeatureStore.get`. --- **`num_tokens`** `int` Sum of attended tokens; used by the consumer for empty / short loss-mask neutralization and by the queue for backpressure. --- **`run_id`** `str` Stable identifier across the whole training run. Same value on every ref of one run so producers and consumers can verify they are talking about the same run before materializing. --- **`sample_id`** `str` Stable identifier within `run_id`. The consumer uses it to partition the stream across DP ranks (RFC §"Consumer-side DP resharding") and to ACK / FAIL the corresponding lease. --- **`schema_version`** `int` Bumped whenever the producer's feature set or tensor layout for `algorithm` changes incompatibly. The consumer uses it as a hard gate on the ref. --- **`store_uri`** `str` `scheme://location` identifying the `FeatureStore` back-end (e.g. `mem://local` for `LocalFeatureStore`). The store URI plus `feature_keys` form the lookup the consumer makes; the store object itself is discovered through the URI at materialization time, not stored on the ref. --- **`target_model_version`** `str` Monotonically increasing identifier of the target-model weights that produced this sample. Used to reject refs from a stale target (relevant once train-with-decode / weight resync lands in PR 4). --- ```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 `FeatureStore.get`. ```python class nemo_automodel.components.speculative.streaming.queue.SampleRefQueue( store: nemo_automodel.components.speculative.streaming.store.FeatureStore, visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout | None = None, high_watermark_bytes: int | None = None, low_watermark_bytes: int | None = None, on_pause: typing.Callable[[StoreHealth], None] | None = None, on_resume: typing.Callable[[StoreHealth], None] | None = None ) ``` Lease / ack / fail queue over `SampleRef`. Thread safety: a single `threading.Lock` protects every list / counter, so a multi-producer / multi-consumer deployment works as long as only one thread at a time calls any one of the methods. **Parameters:** **`store`** `FeatureStore` The data-plane store the consumers will materialize against. The queue reads `FeatureStore.health` for backpressure. --- **`visibility_timeout`** `VisibilityTimeout | None` — default: None How long a leased-but-not-acked ref can live before reclaim. Defaults to 30s; production deployments normally key this off the recipe's per-step budget. --- **`high_watermark_bytes`** `int | None` — default: None Optional resident-byte threshold for pausing. When `None` (default), the queue defers to `StoreHealth.high_watermark_hit` (i.e. the store's own configured threshold). When set, the queue pauses whenever `StoreHealth.resident_bytes >= high_watermark_bytes`. --- **`low_watermark_bytes`** `int | None` — default: None Optional resident-byte threshold for resuming. When `None` (default), the queue defers to `StoreHealth.low_watermark_hit`. When set, the queue resumes only after `StoreHealth.resident_bytes <= low_watermark_bytes`. Must be strictly less than `high_watermark_bytes` so the hysteresis band is non-empty. --- **`on_pause / on_resume`** Optional callbacks fired when the queue transitions high-watermark-paused -> resumed and back. --- **`_active_by_sample`** `dict[str, int] = {}` --- **`_lock`** `= threading.Lock()` --- **`_outstanding`** `dict[int, Lease] = {}` --- **`_pending`** `list[SampleRef] = []` --- **`_pending_seen`** `set[str] = set()` --- **`_put_cv`** `= threading.Condition(self._lock)` --- **`_sample_counters`** `dict[str, int] = {}` --- **`_vt`** `= visibility_timeout or VisibilityTimeout()` --- **`is_closed`** `bool` Whether `close` has been called on this queue. Consumers that pull `acquire` and receive `None` use this to disambiguate "drained, stop" (`is_closed is True`) from "transient empty poll, retry" (`is_closed is False`). Mirrors the Python `queue.Queue` separation between `empty()` and the lifecycle-shutdown signal. --- ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._gc_store() -> None ``` ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_pause( health: nemo_automodel.components.speculative.streaming.store.StoreHealth ) -> bool ``` Whether the producer should pause against `health`. When the queue ctor was given an explicit `high_watermark_bytes`, that threshold wins; otherwise the decision defers to `StoreHealth.high_watermark_hit` (i.e. the store's own configured threshold). ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_resume( health: nemo_automodel.components.speculative.streaming.store.StoreHealth ) -> bool ``` Whether the producer should resume against `health`. When the queue ctor was given an explicit `low_watermark_bytes`, that threshold wins; otherwise the decision defers to `StoreHealth.low_watermark_hit`. Hysteresis is preserved either way: resume crosses the low threshold, pause crosses the high threshold. ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.ack( lease: nemo_automodel.components.speculative.streaming.queue.Lease ) -> None ``` Mark a leased ref as successfully consumed and free its queue slot. Verifies `Lease.lease_id` matches the live outstanding entry for `lease.ref.sample_id`: a stale ACK for a lease that has been reclaimed and re-leased is rejected (logged, ignored) so the new consumer's live lease is not popped by accident. Does NOT touch the store -- the consumer's `FeatureStore.get` return value carries a `StoreHandle` that the consumer must hand to `FeatureStore.release` to drop the tensors. The queue's responsibility ends at "lease no longer held". ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.acquire( poll_interval: float = 0.05 ) -> nemo_automodel.components.speculative.streaming.queue.Lease | None ``` Lease the next ref; returns `None` when nothing is ready. `None` is returned in two situations, which consumers disambiguate with `is_closed`: * `is_closed is True`: the queue has been shut down and is drained. The consumer should stop iterating. * `is_closed is False`: a transient empty poll (the producer is briefly behind). The consumer should retry. The returned `Lease` is the only sanctioned way to access the ref's tensors -- `FeatureStore.get` requires a `SampleRef`, and that ref must come from a lease. The consumer MUST hand back the lease via `ack` (on success) or `fail` (on error) so the queue can reclaim the slot and the store can drop the sample. ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.close() -> None ``` Mark the queue closed; `acquire` drains what remains, then returns `None`. Closing does not discard already-enqueued refs: `acquire` keeps handing out pending refs until they are all leased, and only returns `None` once the queue is closed *and* both pending and outstanding are empty. `is_closed` therefore reports the closed flag, not that the queue is already drained; `FeatureDataLoader` polls it to know when a `None` from `acquire` is terminal. Outstanding leases are left intact: their consumer still owns the tensors, and a leaked `FeatureStore.release` would push the store's residency counter below zero. The store's own `close` is the canonical place to drop residency. ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.fail( lease: nemo_automodel.components.speculative.streaming.queue.Lease ) -> None ``` Return a leased ref to the pending queue, without dropping its tensors. Verifies the lease identity before re-enqueuing: a stale `fail` for a lease that has been reclaimed is a no-op. The ref will be leased again (its `Lease.redelivery_count` increments). Re-delivery is what makes the pipeline fault-tolerant to a transient consumer error -- a permanently bad ref is the consumer's problem (drop it after a bounded retry budget). ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.outstanding_count() -> int ``` ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.pending_count() -> int ``` ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put( ref: nemo_automodel.components.speculative.streaming.refs.SampleRef ) -> None ``` Enqueue `ref` for a future `acquire`. Does not block on backpressure; producers that care should call `put_blocks_until_below` instead, which honors the high/low watermark hysteresis from `FeatureStore.health`. ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put_blocks_until_below( ref: nemo_automodel.components.speculative.streaming.refs.SampleRef, poll_interval: float = 0.05, abort_when: typing.Callable[[], bool] | None = None ) -> None ``` Enqueue `ref`, blocking the producer while the store is over its high watermark. The producer is paused when `_should_pause` returns `True` (resident crossed the high threshold) and only resumed when `_should_resume` returns `True` (resident dropped back below the low threshold). In the band between the two thresholds the producer's existing paused / unpaused state is preserved -- that hysteresis is what prevents flapping when the producer is sitting near the high watermark. **Parameters:** **`ref`** `SampleRef` The reference to enqueue. --- **`poll_interval`** `float` — default: 0.05 Seconds between backpressure checks when paused. Defaults to 50ms -- well below typical step times, well above the cost of a Python-level `FeatureStore.health` call. --- **`abort_when`** `Callable[[], bool] | None` — default: None Optional callable checked on each loop iteration. When it returns `True`, the put aborts with `RuntimeError` so a shutdown signal can unblock a producer waiting on backpressure without closing the queue first. --- **Raises:** * `RuntimeError`: if the queue is closed while the producer is blocked, or if `abort_when` returns `True`. ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.reclaim_expired() -> int ``` Reclaim leases whose `Lease.deadline` has passed. Each reclaimed lease is re-enqueued; `acquire` returns it on a future call with an incremented `Lease.redelivery_count`. Returns the number of leases reclaimed -- a queue that is healthy returns 0 most of the time. ```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 ``` ```python class nemo_automodel.components.speculative.streaming.store.StoreHandle( store: 'FeatureStore', sample_id: str, ref: nemo_automodel.components.speculative.streaming.refs.SampleRef, handle_id: int = _next_handle_id() ) ``` Dataclass Opaque token the consumer must return to `FeatureStore.release`. Each `FeatureStore.get` mints a fresh handle with a unique `handle_id`. Two `get` calls on the same sample return two distinct handles; `FeatureStore.release` matches against `handle_id` so releasing one handle twice (or releasing a stale handle after a sibling has been acquired) cannot decrement a sibling's outstanding count. Holds the producing store, the sample id, the originating `SampleRef`, and the per-`get` handle identity. **`handle_id`** `int = field(default_factory=_next_handle_id)` --- **`ref`** `SampleRef` --- **`sample_id`** `str` --- **`store`** `'FeatureStore'` --- ```python class nemo_automodel.components.speculative.streaming.store.StoreHealth( resident_bytes: int, capacity_bytes: int, sample_count: int, high_watermark_hit: bool, low_watermark_hit: bool ) ``` Dataclass Integers-only snapshot of store residency for backpressure decisions. **`capacity_bytes`** `int` Configured hard cap from `LocalFeatureStore` (PR 3's shared-dir / PR 4's NCCL store report the analogous cap). --- **`high_watermark_hit`** `bool` True iff `resident_bytes >= high_watermark_bytes` on the last `health` call. The queue pauses a producer that sees this transition. --- **`low_watermark_hit`** `bool` True iff `resident_bytes <= low_watermark_bytes` on the last `health` call. The queue resumes a producer that has been paused and now sees this transition. The hysteresis band between the two is what prevents flapping. --- **`resident_bytes`** `int` Bytes currently held in the store across un-acked samples. Compared against `capacity_bytes` for the high/low watermark hysteresis. --- **`sample_count`** `int` Number of un-acked samples currently held. Used for the second cap (sample count) the RFC §"Open questions" Q2 keeps alongside the byte backstop. --- ```python class nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout( seconds: float = 30.0 ) ``` Dataclass How long an unacked `Lease` is allowed to live before reclaim. Any positive value is accepted (sub-second values are useful in tests). Production deployments typically pick something an order of magnitude larger than the recipe's per-step budget so a slow but healthy consumer does not see its leases reclaimed out from under it. **`seconds`** `float = 30.0` --- ```python nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout.__post_init__() -> None ``` ```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 `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. ## Docs - [nemo_automodel.components.speculative.streaming.async_pipeline](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/async_pipeline.md) - [nemo_automodel.components.speculative.streaming.eagle3](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/eagle3.md) - [nemo_automodel.components.speculative.streaming.loader](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/loader.md) - [nemo_automodel.components.speculative.streaming.producer](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/producer.md) - [nemo_automodel.components.speculative.streaming.queue](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/queue.md) - [nemo_automodel.components.speculative.streaming.refs](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/refs.md) - [nemo_automodel.components.speculative.streaming.store](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/store.md) - [nemo_automodel.components.speculative.streaming.stores](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/stores.md) - [nemo_automodel.components.speculative.streaming.stores.local](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/stores/local.md) - [nemo_automodel.components.speculative.streaming.stores.shared_dir](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/speculative/streaming/stores/shared_dir.md)