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

Streaming producer for speculative-decoding draft training.

The :class:`FeatureProducer` wraps a single target-backend forward pass
and ships its output through a :class:`FeatureStore`. The trainer-side
:class:`~nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader`
re-materializes the same data as an
:class:`~nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch`
and yields it to the trainer loop.

The first producer wraps the existing
:class:`~nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel`
so the colocated path's numerical output travels bit-for-bit into the
streaming pipeline. Future producers (out-of-process SGLang, NCCL
remote) plug in behind the same `FeatureProducer` API without
touching the trainer.

## Module Contents

### Classes

| Name                                                                                           | Description                                                                     |
| ---------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`FeatureProducer`](#nemo_automodel-components-speculative-streaming-producer-FeatureProducer) | Run a target backend once and put the supervision into a :class:`FeatureStore`. |

### Functions

| Name                                                                                                 | Description                                                     |
| ---------------------------------------------------------------------------------------------------- | --------------------------------------------------------------- |
| [`_resolve_algorithm`](#nemo_automodel-components-speculative-streaming-producer-_resolve_algorithm) | Pick the producer-side :class:`FeatureAlgorithm` for a backend. |

### Data

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

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

### API

```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 :class:`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
:class:`~nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend`.

---

**`store`** `FeatureStore`

The :class:`FeatureStore` every `produce` call writes
into.

---

**`run_id`** `str`

Stable across the whole training run; mirrored onto every
:class:`SampleRef`.

---

**`algorithm`** `FeatureAlgorithm | None` — default: None

Forced :class:`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 :class:`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 :meth:`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 :meth:`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 :class:`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
nemo_automodel.components.speculative.streaming.producer._resolve_algorithm(
    backend: object
) -> nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm
```

Pick the producer-side :class:`FeatureAlgorithm` for a backend.

Dispatch on the backend's base class, not its class-name string: every
Eagle3 target variant (HF, runner, vLLM, SGLang, remote) subclasses
:class:`Eagle3TargetBackend`, so this survives renames, subclassing, and
wrapping. Adding a new algorithm is a one-line `isinstance` branch here.

```python
nemo_automodel.components.speculative.streaming.producer.__all__ = ['FeatureProducer']
```

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