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

DFlash 2 draft model (Qwen3-style).

DFlash 2 ([https://inco.ai/blog/dflash2/](https://inco.ai/blog/dflash2/)) keeps DFlash's one-pass parallel block
draft -- the block is still predicted in a single non-causal forward -- and adds
two cheap modules on top of it:

* A **two-tap dynamic depthwise convolution** before and after each attention and
  MLP sublayer of every draft layer::

  Conv(x)*t = k*\{t,0} \* x\_t + k\_\{t,1} \* x\_\{t-1}

  Each coefficient is a learned base kernel plus a small correction predicted
  from the current hidden state and shared across `conv_group_size` channels.
  The taps never cross a draft-block boundary, so block position 1 reads block
  position 0 -- the last verified (anchor) token -- and block position 0 reads
  zero padding. This moves the short-range within-block work off attention (which
  goes back to reading the target context) and removes most of DFlash's *suffix
  decay*: prediction quality at the end of the block, for \~3% extra parameters.

* A **pairwise path selector**. DFlash picks every position's top-1 candidate
  independently, so neighbours can disagree (a repeated word, a broken phrase)
  and the block is cut short at verification -- even though the right token is
  usually already in the position's candidate list. The selector keeps each
  position's top `selector_top_k` candidates and scores every adjacent pair in
  one shot::

  S\_t(a, b) = U\_t(b) + \<A(a) \* H(h\_t), B(b)>

  `U_t(b)` is DFlash's own logit for candidate `b`; `A` and `B` are
  rank-`selector_rank` token codebooks matched under a context gate `H(h_t)`
  \-- a low-rank bilinear score over adjacent candidates. Scoring is fully
  parallel; only the final walk over the precomputed scores is sequential, and it
  touches no backbone or LM head.

Both modules are initialised to the identity (zero conv correction, zero
successor codebook), so a freshly constructed DFlash 2 draft starts out
numerically equal to plain DFlash and training moves it away from there.

## Module Contents

### Classes

| Name                                                                                                                     | Description                                                                     |
| ------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------- |
| [`CandidateSelector`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-CandidateSelector)               | Pairwise path selector over each draft position's top-k candidates.             |
| [`GroupedDynamicCausalConv`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-GroupedDynamicCausalConv) | Two-tap dynamic depthwise convolution wrapped around one draft sublayer.        |
| [`Qwen3DFlash2DecoderLayer`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-Qwen3DFlash2DecoderLayer) | A DFlash 2 decoder block: DFlash's block + a two-tap conv around each sublayer. |
| [`Qwen3DFlash2DraftModel`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-Qwen3DFlash2DraftModel)     | DFlash 2 draft model: the DFlash stack plus in-block convs and a path selector. |

### Functions

| Name                                                                                                                       | Description                                                                  |
| -------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`_grouped_dynamic_convolve`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-_grouped_dynamic_convolve) | Depthwise convolution over draft-block positions with content-adaptive taps. |
| [`dflash2_rejection_sample`](#nemo_automodel-components-speculative-dflash-draft_qwen3_dflash2-dflash2_rejection_sample)   | Accept a prefix of the drafted block and resample the first rejected token.  |

### API

```python
class nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.CandidateSelector(
    vocab_size: int,
    hidden_size: int,
    rank: int,
    top_k: int
)
```

**Bases:** `Module`

Pairwise path selector over each draft position's top-k candidates.

Scores `S_t(a, b) = U_t(b) + &lt;A(a) * H(h_t), B(b)&gt;` for predecessor token
`a` and candidate `b`: DFlash's own logit plus a low-rank bilinear match
between the two tokens' codebook embeddings, gated by the draft hidden state.

The codebooks are plain `[vocab, rank]` parameters rather than
`nn.Embedding` modules so the saved keys are `candidate_selector.*_codebook`
\-- the names the published DFlash 2 drafters and the serving runtimes use --
instead of `candidate_selector.*_codebook.weight`.

**Parameters:**

**`vocab_size`** `int`

Size of the (shared with the target) token vocabulary.

---

**`hidden_size`** `int`

Channel count of the draft hidden states.

---

**`rank`** `int`

Codebook / gate width (`selector_rank`; 256 in DFlash 2).

---

**`top_k`** `int`

Candidates kept per position (`selector_top_k`; 16 in DFlash 2).

---

**`hidden_projection`** `= nn.Linear(hidden_size, rank, bias=False)`

---

**`predecessor_codebook`** `= nn.Parameter(torch.empty(vocab_size, rank))`

---

**`successor_codebook`** `= nn.Parameter(torch.empty(vocab_size, rank))`

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.CandidateSelector.pair_scores(
    hidden: torch.Tensor,
    unary: torch.Tensor,
    candidate_ids: torch.Tensor,
    predecessor_ids: torch.Tensor
) -> torch.Tensor
```

Score every (predecessor, candidate) pair for a batch of draft positions.

Fully parallel: no position depends on another's score. Training calls this
once with the ground-truth predecessors; the decode-time walk in
:meth:`walk` calls it one position at a time with the token it just picked.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[..., hidden]; the draft hidden state at each
scored position, with arbitrary leading dimensions.

---

**`unary`** `torch.Tensor`

Tensor of shape \[..., candidates]; `U_t(b)`, the draft logit of
each candidate, with the same leading dimensions as `hidden`.

---

**`candidate_ids`** `torch.Tensor`

Long tensor of shape \[..., candidates]; the candidate
token ids, with the same leading dimensions as `hidden`.

---

**`predecessor_ids`** `torch.Tensor`

Long tensor of shape \[...]; the token preceding each
scored position, with the same leading dimensions as `hidden`.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., candidates] holding `S_t(a, b)`.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.CandidateSelector.reset_parameters() -> None
```

Reset the selector to a no-op: every score collapses to the draft logit.

`S_t` is bilinear in the two codebooks, so collapsing it needs one factor
at zero. Zeroing the successor codebook does that; it still receives
gradient immediately, while the predecessor codebook and the context gate
multiply it and therefore only start training on the second step.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.CandidateSelector.walk(
    hidden: torch.Tensor,
    logits: torch.Tensor,
    anchor_ids: torch.Tensor,
    temperature: float
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Trace one coherent path through the per-position candidate lists.

Greedy (`temperature == 0`) follows the best successor at each step;
otherwise the step is sampled from the softmax over the candidate scores,
and the returned per-step distribution is the draft proposal `q` that
:func:`dflash2_rejection_sample` needs to stay lossless.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, draft, hidden]; the draft hidden states
of the block's predicted positions.

---

**`logits`** `torch.Tensor`

Tensor of shape \[batch, draft, vocab]; the draft logits at those
positions.

---

**`anchor_ids`** `torch.Tensor`

Long tensor of shape \[batch]; the last verified token, i.e.
the predecessor of draft position 0.

---

**`temperature`** `float`

Sampling temperature; `0` selects greedily.

---

**Returns:** `torch.Tensor`

Tuple `(path, candidate_ids, draft_probs)`: `path` is a Long tensor

```python
class nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.GroupedDynamicCausalConv(
    hidden_size: int,
    kernel_size: int,
    group_size: int
)
```

**Bases:** `Module`

Two-tap dynamic depthwise convolution wrapped around one draft sublayer.

One instance covers the convolution *before* a sublayer (:meth:`prepare`) and
the one *after* it (:meth:`finish`). Both sets of dynamic coefficients are
predicted from the sublayer's input, so :meth:`prepare` returns the
coefficients :meth:`finish` needs and the projection runs once per sublayer.

**Parameters:**

**`hidden_size`** `int`

Channel count of the draft hidden states.

---

**`kernel_size`** `int`

Number of taps; DFlash 2 uses 2 (self and predecessor).

---

**`group_size`** `int`

Channels sharing one dynamic correction (16 in DFlash 2).

---

**`base_kernel`**

---

**`kernel_projection`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.GroupedDynamicCausalConv.finish(
    hidden: torch.Tensor,
    dynamic: torch.Tensor,
    block_size: int
) -> torch.Tensor
```

Apply the post-sublayer convolution using the coefficients from :meth:`prepare`.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; the sublayer output.

---

**`dynamic`** `torch.Tensor`

Tensor of shape \[batch, sequence, kernel, groups]; the second
half of :meth:`prepare`'s coefficients.

---

**`block_size`** `int`

Draft-block length; taps never cross a block boundary.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.GroupedDynamicCausalConv.prepare(
    hidden: torch.Tensor,
    block_size: int
) -> tuple[torch.Tensor, torch.Tensor]
```

Apply the pre-sublayer convolution and emit the post-sublayer coefficients.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; the sublayer input.

---

**`block_size`** `int`

Draft-block length; taps never cross a block boundary.

---

**Returns:** `torch.Tensor`

Tuple of `(convolved, dynamic)`: `convolved` is a Tensor of shape

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.GroupedDynamicCausalConv.reset_parameters() -> None
```

Reset the convolution to the identity: unit self-tap, no correction.

```python
class nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DecoderLayer(
    config: transformers.models.qwen3.configuration_qwen3.Qwen3Config,
    layer_idx: int
)
```

**Bases:** [Qwen3DFlashDecoderLayer](/nemo-automodel/nemo_automodel/components/speculative/dflash/draft_qwen3#nemo_automodel-components-speculative-dflash-draft_qwen3-Qwen3DFlashDecoderLayer)

A DFlash 2 decoder block: DFlash's block + a two-tap conv around each sublayer.

**`attention_conv`**

---

**`mlp_conv`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DecoderLayer.forward(
    target_hidden: torch.Tensor | None = None,
    hidden_states: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_value: transformers.cache_utils.Cache | None = None,
    use_cache: bool | None = False,
    cache_position: torch.LongTensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    conv_block_size: int = 1,
    kwargs = {}
) -> torch.Tensor
```

Run the block, convolving the input and output of each sublayer.

**Parameters:**

**`target_hidden`** `torch.Tensor | None` — default: None

Tensor of shape \[batch, context, hidden]; the projected
target-model context the attention keys/values are extended with.

---

**`hidden_states`** `torch.Tensor | None` — default: None

Tensor of shape \[batch, draft, hidden]; the draft
(noise-block) positions, `draft` being a whole number of
`conv_block_size`-long blocks.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Attention mask over \[batch, 1, draft, context + draft],
a flex `BlockMask`, or `None`.

---

**`position_ids`** `torch.LongTensor | None` — default: None

Long tensor of shape \[batch, context + draft].

---

**`past_key_value`** `Cache | None` — default: None

Draft KV cache, or `None`.

---

**`use_cache`** `bool | None` — default: False

Whether to write into `past_key_value`.

---

**`cache_position`** `torch.LongTensor | None` — default: None

Long tensor of shape \[draft], or `None`.

---

**`position_embeddings`** `tuple[torch.Tensor, torch.Tensor] | None` — default: None

Tuple of rotary `(cos, sin)` tensors of shape
\[batch, context + draft, head\_dim].

---

**`conv_block_size`** `int` — default: 1

Draft-block length; the convolutions' predecessor tap
never crosses a block boundary.

---

**`**kwargs`** — default: \{}

Forwarded to the attention implementation.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, draft, hidden].

```python
class nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel(
    config
)
```

**Bases:** [Qwen3DFlashDraftModel](/nemo-automodel/nemo_automodel/components/speculative/dflash/draft_qwen3#nemo_automodel-components-speculative-dflash-draft_qwen3-Qwen3DFlashDraftModel)

DFlash 2 draft model: the DFlash stack plus in-block convs and a path selector.

**`_no_split_modules`** `= ['Qwen3DFlash2DecoderLayer']`

---

**`candidate_selector`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel.forward(
    position_ids: torch.LongTensor,
    attention_mask: torch.Tensor | None = None,
    noise_embedding: torch.Tensor | None = None,
    target_hidden: torch.Tensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    use_cache: bool = False,
    conv_block_size: int | None = None,
    kwargs = {}
) -> torch.Tensor
```

Run the DFlash 2 draft stack over `[context | noise-block]`.

**Parameters:**

**`position_ids`** `torch.LongTensor`

Long tensor of shape \[batch, context + draft].

---

**`attention_mask`** `torch.Tensor | None` — default: None

Attention mask over \[batch, 1, draft, context + draft],
a flex `BlockMask`, or `None`.

---

**`noise_embedding`** `torch.Tensor | None` — default: None

Tensor of shape \[batch, draft, hidden]; the embedded
`[anchor, MASK, ...]` blocks laid end to end.

---

**`target_hidden`** `torch.Tensor | None` — default: None

Tensor of shape \[batch, context, layers \* hidden]; the
concatenated target-model context features.

---

**`past_key_values`** `Cache | None` — default: None

Draft KV cache, or `None`.

---

**`use_cache`** `bool` — default: False

Whether to write into `past_key_values`.

---

**`conv_block_size`** `int | None` — default: None

Draft-block length for the in-block convolutions; see
:meth:`resolve_conv_block_size` for how `None` is resolved.

---

**`**kwargs`** — default: \{}

Forwarded to the attention implementation.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, draft, hidden]; the normalised draft hidden

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel.resolve_conv_block_size(
    query_len: int,
    conv_block_size: int | None
) -> int
```

Resolve the block length the in-block convolutions must not reach across.

**Parameters:**

**`query_len`** `int`

Number of draft (noise-block) query positions in this call.

---

**`conv_block_size`** `int | None`

Explicit block length, or `None` to infer it.

---

**Returns:** `int`

The block length to convolve within. `None` resolves to

**Raises:**

* `ValueError`: If `conv_block_size` does not divide `query_len`.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel.spec_generate(
    target: torch.nn.Module,
    input_ids: torch.LongTensor,
    max_new_tokens: int,
    stop_token_ids: list[int] | None,
    temperature: float
) -> torch.LongTensor
```

Block-parallel speculative decoding with pairwise path selection.

Each cycle drafts one block in a single draft forward, walks the selector
over the per-position candidates to pick a coherent path, and verifies the
whole block with one target forward. `temperature == 0` accepts the
longest exact-match prefix; `temperature &gt; 0` accepts via rejection
sampling, so the emitted tokens follow the target's own distribution.

**Parameters:**

**`target`** `nn.Module`

The frozen verifier; must expose `model.embed_tokens`,
`lm_head`, and an HF-style forward with `output_hidden_states`.

---

**`input_ids`** `torch.LongTensor`

Long tensor of shape \[1, prompt].

---

**`max_new_tokens`** `int`

Maximum number of tokens to generate.

---

**`stop_token_ids`** `list[int] | None`

Token ids that end generation, or `None`.

---

**`temperature`** `float`

Sampling temperature; `0` decodes greedily.

---

**Returns:** `torch.LongTensor`

Long tensor of shape \[1, prompt + generated] containing the prompt

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2._grouped_dynamic_convolve(
    hidden: torch.Tensor,
    dynamic: torch.Tensor,
    base_kernel: torch.Tensor,
    group_size: int,
    block_size: int
) -> torch.Tensor
```

Depthwise convolution over draft-block positions with content-adaptive taps.

Tap `offset` reads position `t - offset` of the *same* draft block; the
leading `offset` positions of every block read zero padding instead of the
previous block's tail, which is what keeps the packed `blocks * block_size`
training layout equivalent to drafting one block at a time.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; `sequence` is a whole
number of `block_size`-long draft blocks laid end to end.

---

**`dynamic`** `torch.Tensor`

Tensor of shape \[batch, sequence, kernel, groups]; the per-position
correction added to `base_kernel`, shared by the `group_size`
channels of each group.

---

**`base_kernel`** `torch.Tensor`

Tensor of shape \[kernel, hidden]; the learned base taps.

---

**`group_size`** `int`

Number of channels sharing one dynamic correction.

---

**`block_size`** `int`

Draft-block length; taps never cross a block boundary.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden]; a fresh tensor that neither

```python
nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.dflash2_rejection_sample(
    draft_tokens: torch.Tensor,
    target_probs: torch.Tensor,
    draft_probs: torch.Tensor,
    candidate_ids: torch.Tensor
) -> tuple[int, torch.Tensor]
```

Accept a prefix of the drafted block and resample the first rejected token.

Standard speculative-decoding rejection sampling, specialised to a proposal
supported on the selector's candidate set: `q(token)` is read out of
`draft_probs` by matching `candidate_ids`, and the residual
`clamp(p - q, 0)` is formed by scattering `-q` into the full-vocabulary
target distribution. The accepted tokens plus the resampled one are therefore
distributed exactly as the target's own samples.

**Parameters:**

**`draft_tokens`** `torch.Tensor`

Long tensor of shape \[1, draft]; the selector's path.

---

**`target_probs`** `torch.Tensor`

Tensor of shape \[1, block, vocab]; the verifier's
next-token distribution at each block position, where
`block == draft + 1`.

---

**`draft_probs`** `torch.Tensor`

Tensor of shape \[1, draft, candidates]; the proposal mass on
each candidate.

---

**`candidate_ids`** `torch.Tensor`

Long tensor of shape \[1, draft, candidates]; the candidate
token ids the proposal is supported on.

---

**Returns:** `int`

Tuple `(accepted, bonus)`: `accepted` is the number of drafted tokens