nemo_automodel.components.speculative.dspark.common

Module Contents

Classes

Name	Description
`AcceptRatePredictor`	-
`DSparkForwardOutput`	Outputs for one DSpark training forward.

Functions

Name	Description
`build_anchor_candidate_mask`	-
`build_eval_mask`	-
`context_doc_ids`	Per-context-token document id `[B, S]` from packed `seq_lens` `[B, max_docs]`.
`create_noise_embed`	-
`create_position_ids`	Position ids for the parallel draft blocks (`base position + offset`).
`extract_context_feature`	-
`pin_rope_inv_freq_fp32`	Keep a RoPE module’s `inv_freq` buffers in fp32 after a dtype cast.
`sample_anchor_positions`	-
`validate_target_layer_ids`	-

Data

__all__

API

class nemo_automodel.components.speculative.dspark.common.AcceptRatePredictor(
    input_dim: int
)

Bases: Module

proj

= nn.Linear(int(input_dim), 1)

nemo_automodel.components.speculative.dspark.common.AcceptRatePredictor.forward(
    features
)

class nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput(
    draft_logits: torch.Tensor,
    target_ids: torch.Tensor,
    eval_mask: torch.Tensor,
    block_keep_mask: torch.Tensor,
    confidence_pred: typing.Optional[torch.Tensor] = None,
    aligned_target_logits: typing.Optional[torch.Tensor] = None
)

Dataclass

Outputs for one DSpark training forward.

The sampler keeps anchors whose first draft target is enabled by loss_mask. Later slots are supervised only while they remain inside seq_len and form a contiguous enabled prefix. Dummy anchors can still appear when a sample has too few valid anchors; they are masked out by block_keep_mask and eval_mask.

aligned_target_logits

Optional[Tensor] = None

block_keep_mask

Tensor

confidence_pred

Optional[Tensor] = None

draft_logits

Tensor

eval_mask

Tensor

target_ids

Tensor

nemo_automodel.components.speculative.dspark.common.build_anchor_candidate_mask(
    seq_len: int,
    loss_mask: torch.Tensor,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.common.build_eval_mask(
    seq_len: int,
    loss_mask: torch.Tensor,
    label_indices: torch.Tensor,
    safe_label_indices: torch.Tensor,
    block_keep_mask: torch.Tensor,
    doc_remaining: typing.Optional[torch.Tensor] = None,
    anchor_positions: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.common.context_doc_ids(
    seq_lens: torch.Tensor,
    seq_len: int,
    device: torch.device
) -> torch.Tensor

Per-context-token document id [B, S] from packed seq_lens [B, max_docs].

Mirrors the doc_id construction in build_block_causal_additive_mask: a token’s id is the number of document boundaries at or before its position, so 0-length padding entries never split a real document.

nemo_automodel.components.speculative.dspark.common.create_noise_embed(
    embed_tokens: torch.nn.Module,
    input_ids: torch.Tensor,
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    mask_token_id: int,
    block_size: int
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.common.create_position_ids(
    anchor_positions: torch.Tensor,
    block_size: int,
    context_position_ids: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

Position ids for the parallel draft blocks (base position + offset).

Without packing the anchor’s position equals its row index, so the block positions are anchor + offset. Under packing context_position_ids [B, S] holds per-document reset positions, so the block’s base position is gathered from it at the anchor to keep the draft’s RoPE phase document-local.

nemo_automodel.components.speculative.dspark.common.extract_context_feature(
    hidden_states,
    layer_ids
)

nemo_automodel.components.speculative.dspark.common.pin_rope_inv_freq_fp32(
    rotary_emb: typing.Optional[torch.nn.Module]
) -> None

Keep a RoPE module’s inv_freq buffers in fp32 after a dtype cast.

module.to(bfloat16) (the training build path) rounds the rotary inv_freq buffers to bf16. The rounded frequencies dephase with absolute position, so the train/inference RoPE diverges (worse with longer context) and draft acceptance erodes, while the serving runtime keeps an fp32 RoPE cache. A bf16 round-trip cannot be undone by upcasting, so recompute fresh fp32 frequencies from a fresh rotary module built off the same config (the same values HF derives on the fp32 paths) and copy them back in.

Works for both the single-buffer layout (inv_freq / original_inv_freq, e.g. Qwen3) and the per-layer-type layout where each frequency buffer is named <layer_type>_inv_freq (e.g. Gemma4). No-op when every frequency buffer is already fp32 or on a meta device.

nemo_automodel.components.speculative.dspark.common.sample_anchor_positions(
    seq_len: int,
    loss_mask: torch.Tensor,
    num_anchors: int,
    device: torch.device,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.common.validate_target_layer_ids(
    layer_ids,
    num_target_layers: int
)

nemo_automodel.components.speculative.dspark.common.__all__ = ['DSparkForwardOutput', 'AcceptRatePredictor', 'extract_context_feature', 'valid...