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

1. **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`.
2. **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`](#nemo_automodel-components-speculative-dflash-jetspec_core-JetSpecStepMetrics)     | Per-step training outputs for the JetSpec draft.                                     |
| [`JetSpecTrainerModule`](#nemo_automodel-components-speculative-dflash-jetspec_core-JetSpecTrainerModule) | JetSpec online training wrapper: causal parallel drafting + forward-KL distillation. |

### Data

[`_IGNORE_INDEX`](#nemo_automodel-components-speculative-dflash-jetspec_core-_IGNORE_INDEX)

### API

```python
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.

```python
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](/nemo-automodel/nemo_automodel/components/speculative/dflash/core#nemo_automodel-components-speculative-dflash-core-DFlashTrainerModule)

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

```python
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]`.

```python
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.

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