nemo_automodel.components.speculative.streaming.eagle3

EAGLE-3 schema for the streaming data plane.

Mirrors the EAGLE-3 row of the speculative-decoding train/inference disaggregation RFC’s “Feature schema per algorithm” table:

Required feature keys	Supervision
aux_hidden_states (3 aux concat, H*3),	logits [B,S,V] OR target_probs + position_mask
input_ids, attention_mask, loss_mask	(draft-vocab; mutually exclusive)

“Exactly one supervision encoding” is the rule. The colocated target backend ships full logits; the remote target backend ships target_probs + position_mask so the wire never carries a full-vocab tensor.

DFlash and DSpark schemas land alongside their producer / loader.

This is a greenfield API: bump EAGLE3_SCHEMA_VERSION when the feature layout changes; consumers reject mismatched refs rather than migrating older samples.

Module Contents

Functions

Name	Description
`eagle3_logits_feature_specs`	Build the per-feature :class:`FeatureSpec` map for a colocated encoder batch.
`eagle3_logits_tensors`	Pack an EAGLE-3 colocated-path encoder’s outputs into the producer’s tensor dict.
`validate_eagle3_packing_inputs`	Require all packing metadata together or none at all.
`validate_eagle3_ref`	Verify `ref` matches the EAGLE-3 schema before the consumer materializes.

Data

EAGLE3_CORE_FEATURES

EAGLE3_DRAFT_VOCAB_SUPERVISION

EAGLE3_LOGITS_SUPERVISION

EAGLE3_PACKING_FEATURES

EAGLE3_SCHEMA_VERSION

EAGLE3_SUPERVISION_ENCODINGS

__all__

logger

API

nemo_automodel.components.speculative.streaming.eagle3.eagle3_logits_feature_specs(
    aux_hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    logits: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> dict[str, nemo_automodel.components.speculative.streaming.refs.FeatureSpec]

Build the per-feature :class:FeatureSpec map for a colocated encoder batch.

Returns: dict[str, FeatureSpec]

A dict[str, FeatureSpec] keyed by feature name with the

nemo_automodel.components.speculative.streaming.eagle3.eagle3_logits_tensors(
    aux_hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    logits: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> dict[str, torch.Tensor]

Pack an EAGLE-3 colocated-path encoder’s outputs into the producer’s tensor dict.

Parameters:

aux_hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden * num_aux_layers]. num_aux_layers is the count of distinct layers the target’s forward hooks captured (3 for the EAGLE-3 default recipe, concat’d along the last axis).

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.

logits

torch.Tensor

Tensor of shape [batch, sequence, vocab]. Per the colocated path this is the target’s full LM-head output.

position_ids

torch.Tensor | NoneDefaults to None

Optional [batch, sequence] per-document positions when sequence packing is enabled.

seq_lens

torch.Tensor | NoneDefaults to None

Optional [batch, max_docs] packed document lengths.

doc_remaining

torch.Tensor | NoneDefaults to None

Optional [batch, sequence] cross-document TTT gate.

Returns: dict[str, torch.Tensor]

A dict[str, torch.Tensor] keyed by

nemo_automodel.components.speculative.streaming.eagle3.validate_eagle3_packing_inputs(
    position_ids: torch.Tensor | None,
    seq_lens: torch.Tensor | None,
    doc_remaining: torch.Tensor | None
) -> None

Require all packing metadata together or none at all.

nemo_automodel.components.speculative.streaming.eagle3.validate_eagle3_ref(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef
) -> None

Verify ref matches the EAGLE-3 schema before the consumer materializes.

Parameters:

ref

SampleRef

The tensor-free reference carried by the queue’s lease.

Raises:

ValueError: if ref.algorithm is not :data:FeatureAlgorithm.EAGLE3, if any core feature is missing from the ref, or if the supervision encoding is neither logits alone nor the target_probs + position_mask pair.

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_CORE_FEATURES: tuple[str, ...] = ('aux_hidden_states', 'input_ids', 'attention_mask', 'loss_mask')

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_DRAFT_VOCAB_SUPERVISION: tuple[str, ...] = ('target_probs', 'position_mask')

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_LOGITS_SUPERVISION: tuple[str, ...] = ('logits',)

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_PACKING_FEATURES: tuple[str, ...] = ('position_ids', 'seq_lens', 'doc_remaining')

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_SCHEMA_VERSION = 1

nemo_automodel.components.speculative.streaming.eagle3.EAGLE3_SUPERVISION_ENCODINGS: tuple[str, ...] = EAGLE3_CORE_FEATURES + EAGLE3_LOGITS_SUPERVISION

nemo_automodel.components.speculative.streaming.eagle3.__all__ = ['EAGLE3_CORE_FEATURES', 'EAGLE3_DRAFT_VOCAB_SUPERVISION', 'EAGLE3_LOGITS_SUPERV...

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