nemo_automodel.components.speculative.dflash.jetspec_core

JetSpec online training wrapper.

JetSpec (arXiv:2606.18394, “Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting”) reuses the DFlash parallel draft backbone but makes two changes that turn the block-parallel draft into a causal parallel tree drafter:

Block-causal attention (paper §2.2). DFlash drafts a block bidirectionally, so each node’s distribution is branch-agnostic and the constructed tree can be internally inconsistent. JetSpec masks the in-block attention causally — a query at within-block offset i attends only to offsets j <= i — so each branch is conditioned on its own ancestor prefix and the draft factorization mirrors the target’s autoregressive order. This is the causal=True path of :func:~nemo_automodel.components.attention.dflash_mask.create_dflash_block_mask.
Forward-KL distillation (paper §2.3, Eq. 8-9). Instead of DFlash’s hard-label decay-weighted CE, JetSpec matches the target model’s per-position soft distribution with a temperature-scaled forward-KL objective (:class:~nemo_automodel.components.loss.kd_loss.KDLoss), so the draft preserves the teacher’s relative preferences across plausible continuations.

JetSpecTrainerModule subclasses :class:DFlashTrainerModule and reuses its anchor sampling, [anchor, MASK, ...] noise-block construction, and absolute position ids; only the attention mask (causal) and the loss (forward-KL against teacher logits) are JetSpec-specific. The draft model itself is the unmodified Qwen3DFlashDraftModel — the mask is supplied by this wrapper, so no new draft architecture is needed.

Module Contents

Classes

Name	Description
`JetSpecStepMetrics`	Per-step training outputs for the JetSpec draft.
`JetSpecTrainerModule`	JetSpec online training wrapper: causal parallel drafting + forward-KL distillation.

Data

_IGNORE_INDEX

API

class nemo_automodel.components.speculative.dflash.jetspec_core.JetSpecStepMetrics(
    loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor,
    accept_len: torch.Tensor
)

Dataclass

Per-step training outputs for the JetSpec draft.

loss/accuracy/valid_tokens mirror DFlashStepMetrics so the shared DFlash training loop consumes them unchanged. accept_len is the expected accepted-prefix length per block (a greedy acceptance-length / tau proxy), which is the headline quantity for speculative decoding — far more informative than the depth-averaged token accuracy.

accept_len

Tensor

accuracy

Tensor

loss

Tensor

valid_tokens

Tensor

class nemo_automodel.components.speculative.dflash.jetspec_core.JetSpecTrainerModule(
    draft_model: nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel,
    target_lm_head: torch.nn.Module,
    target_embed_tokens: torch.nn.Module,
    mask_token_id: int,
    block_size: int = 16,
    attention_backend: str = 'flex_attention',
    num_anchors: int = 512,
    kd_temperature: float = 1.0,
    kd_chunk_size: int = 0
)

Bases: DFlashTrainerModule

JetSpec online training wrapper: causal parallel drafting + forward-KL distillation.

kd_chunk_size

= int(kd_chunk_size)

kd_loss_fn

kd_temperature

= float(kd_temperature)

nemo_automodel.components.speculative.dflash.jetspec_core.JetSpecTrainerModule._gather_teacher_logits(
    target_logits: torch.Tensor,
    label_indices: torch.Tensor,
    seq_len: int
) -> torch.Tensor

Teacher distribution for each predicted position, gathered from the target logits.

Block offset k (k = 1..bs-1) predicts the token at sequence position anchor + k; the target model’s autoregressive distribution for that token, on the ground-truth prefix, is its logits at position anchor + k - 1. Those source positions are label_indices[..., :-1] (anchor + 0 .. anchor + bs-2). Returns [B, N*(bs-1), V].

nemo_automodel.components.speculative.dflash.jetspec_core.JetSpecTrainerModule.forward(
    input_ids: torch.Tensor,
    hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_logits: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dflash.jetspec_core.JetSpecStepMetrics

Causal parallel block-wise forward with a forward-KL distillation loss.

target_logits is the frozen target’s full-vocab logits [B, S, V] (captured by HFDFlashTargetModel(capture_logits=True)); it supplies the teacher distribution for every supervised draft position. Under sequence packing (position_ids / seq_lens / doc_remaining, see DFlashTrainerModule.forward) the target runs block-causal, so the teacher logits gathered below are document-local, and the anchor sampling keeps every gathered teacher position (up to anchor + block_size - 2) inside the anchor’s document.

nemo_automodel.components.speculative.dflash.jetspec_core._IGNORE_INDEX = -100