> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

DFlash draft model (Qwen3-style).

Ported from SpecForge's `specforge/modeling/draft/dflash.py`. DFlash drafts a
whole block of `block_size` tokens in parallel: the block's first position
holds the real anchor token and the rest are `MASK` tokens, and the draft
predicts the whole block in a single non-causal forward conditioned on the
target model's context hidden states.

The draft attention is therefore **not causal** -- a draft block's queries
attend to (a) the projected target-hidden context strictly before its anchor and
(b) bidirectionally to the other (noise) tokens of the same block. The attention
mask that enforces this is built by the trainer wrapper in
`nemo_automodel.components.speculative.dflash.core`.

## Module Contents

### Classes

| Name                                                                                                           | Description                                                                   |
| -------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`Qwen3DFlashAttention`](#nemo_automodel-components-speculative-dflash-draft_qwen3-Qwen3DFlashAttention)       | Non-causal attention whose keys/values are `[context \| noise-block]`.        |
| [`Qwen3DFlashDecoderLayer`](#nemo_automodel-components-speculative-dflash-draft_qwen3-Qwen3DFlashDecoderLayer) | A DFlash decoder block: non-causal attention over `[context \| noise]` + MLP. |
| [`Qwen3DFlashDraftModel`](#nemo_automodel-components-speculative-dflash-draft_qwen3-Qwen3DFlashDraftModel)     | DFlash draft model: a small non-causal Qwen3 stack over `[context \| noise]`. |

### Functions

| Name                                                                                                                           | Description                                                                       |
| ------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------- |
| [`_sliding_window_mask`](#nemo_automodel-components-speculative-dflash-draft_qwen3-_sliding_window_mask)                       | Band mask keeping keys within `sliding_window` of each query's position.          |
| [`apply_rotary_pos_emb`](#nemo_automodel-components-speculative-dflash-draft_qwen3-apply_rotary_pos_emb)                       | Apply RoPE where queries (draft block) are a suffix of the key positions.         |
| [`assert_target_supports_rollback`](#nemo_automodel-components-speculative-dflash-draft_qwen3-assert_target_supports_rollback) | Reject targets whose cache cannot be rewound after a rejected block.              |
| [`build_qwen3_dflash_draft_config`](#nemo_automodel-components-speculative-dflash-draft_qwen3-build_qwen3_dflash_draft_config) | Build the DFlash draft config for a Qwen3-shaped target.                          |
| [`build_target_layer_ids`](#nemo_automodel-components-speculative-dflash-draft_qwen3-build_target_layer_ids)                   | Pick `num_draft_layers` target layers spread across the target's depth.           |
| [`extract_context_feature`](#nemo_automodel-components-speculative-dflash-draft_qwen3-extract_context_feature)                 | Concatenate the selected target layers' hidden states along the feature dim.      |
| [`resolve_output_head`](#nemo_automodel-components-speculative-dflash-draft_qwen3-resolve_output_head)                         | Return the target's output projection.                                            |
| [`sample`](#nemo_automodel-components-speculative-dflash-draft_qwen3-sample)                                                   | Greedy (`temperature == 0`) or truncated-distribution sampling over the last dim. |
| [`sampling_probs`](#nemo_automodel-components-speculative-dflash-draft_qwen3-sampling_probs)                                   | Full-vocabulary sampling distribution after temperature, top-k, and top-p.        |
| [`validate_sampling`](#nemo_automodel-components-speculative-dflash-draft_qwen3-validate_sampling)                             | Reject sampling parameters the decode path cannot honor.                          |

### Data

[`GREEDY_TEMPERATURE_EPS`](#nemo_automodel-components-speculative-dflash-draft_qwen3-GREEDY_TEMPERATURE_EPS)

[`_REWINDABLE_LAYER_TYPES`](#nemo_automodel-components-speculative-dflash-draft_qwen3-_REWINDABLE_LAYER_TYPES)

### API

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

**Bases:** `Module`

Non-causal attention whose keys/values are `[context | noise-block]`.

Queries come from the draft (noise) tokens only; keys and values are the
concatenation of the projected target-hidden context and the noise tokens.
The bidirectional/block structure is supplied entirely by `attention_mask`.

**`attention_dropout`** `= config.attention_dropout`

---

**`head_dim`**

---

**`k_norm`**

---

**`k_proj`**

---

**`num_key_value_groups`**

---

**`o_proj`**

---

**`q_norm`**

---

**`q_proj`**

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`sliding_window`**

---

**`v_proj`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden: torch.Tensor,
    position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs = {}
) -> typing.Tuple[torch.Tensor, torch.Tensor | None]
```

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

**Bases:** `GradientCheckpointingLayer`

A DFlash decoder block: non-causal attention over `[context | noise]` + MLP.

**`hidden_size`** `= config.hidden_size`

---

**`input_layernorm`**

---

**`mlp`** `= Qwen3MLP(config)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDecoderLayer.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: typing.Tuple[torch.Tensor, torch.Tensor] | None = None,
    kwargs = {}
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel(
    config
)
```

**Bases:** `Qwen3PreTrainedModel`

DFlash draft model: a small non-causal Qwen3 stack over `[context | noise]`.

**`_no_split_modules`** `= ['Qwen3DFlashDecoderLayer']`

---

**`block_size`**

---

**`decoder_layer_cls`** `type[Qwen3DFlashDecoderLayer] = Qwen3DFlashDecoderLayer`

---

**`emb_dim`** `= dflash_config['emb_dim']`

---

**`embed_proj`**

---

**`fc`**

---

**`final_logit_softcapping`** `= None if softcap is None else float(softcap)`

---

**`gru_hidden_dim`** `= dflash_config['gru_hidden_dim']`

---

**`hidden_norm`**

---

**`input_embedding_scale`**

---

**`layers`**

---

**`mask_token_id`** `= dflash_config.get('mask_token_id', None)`

---

**`norm`**

---

**`prefix_gru`**

---

**`projector_type`** `= dflash_config.get('projector_type', None)`

---

**`pure_draft_prefix_len`** `= dflash_config.get('pure_draft_prefix_len', 0)`

---

**`rotary_emb`** `= Qwen3RotaryEmbedding(config)`

---

**`shift_label`** `= dflash_config.get('shift_label', False)`

---

**`target_layer_ids`**

---

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel._apply(
    fn,
    recurse = True
)
```

Keep the RoPE `inv_freq` buffer in fp32 across dtype casts.

`Qwen3RotaryEmbedding` computes the rotary angles in fp32 but reads the
frequencies from a stored `inv_freq` buffer. `model.to(bfloat16)` -- the
training build path -- rounds that buffer to bf16, whereas the serving
runtime (SGLang keeps an fp32 RoPE cache) and HF's `from_pretrained` reload
keep it in fp32. The resulting train/inference RoPE mismatch grows with
absolute position (the bf16 frequencies dephase) and erodes draft
acceptance, so `inv_freq` must stay fp32 on both the training and reload
paths. A bf16 round-trip cannot be undone by upcasting, so when a cast
rounds the buffer we recompute fresh fp32 frequencies from the rotary
config (the same values HF derives on the fp32 paths) instead of upcasting
the corrupted ones.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel.compute_logits(
    hidden: torch.Tensor,
    output_head: torch.nn.Module
) -> torch.Tensor
```

Project draft hidden states to logits, applying the target's output transform.

Some targets do not read their `lm_head` output raw: Muse Glimmer scales
it by `output_multiplier` and squashes it through
`final_logit_softcapping`. The draft is trained against, and verified
by, those transformed logits, so both training and decoding have to apply
them -- otherwise the draft learns one distribution and is served another.
Both fields live in `dflash_config`; absent (every Qwen3-family target),
this is just the head.

**Parameters:**

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

Tensor of shape \[..., hidden]; draft hidden states with
arbitrary leading dimensions.

---

**`output_head`** `nn.Module`

The frozen target's output projection.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., vocab].

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel.embed_noise_block(
    target: torch.nn.Module,
    block_ids: torch.LongTensor
) -> torch.Tensor
```

Embed a draft block with the target's (frozen, shared) input embedding.

Applies `input_embedding_scale` from `dflash_config` -- 1.0, hence a
no-op, for every Qwen3-family target. Unlike the reference, this calls the
embedding *module* rather than indexing its weight table directly: under
tensor parallelism the target's `embed_tokens` is vocab-parallel and only
the module call carries the DTensor plan.

**Parameters:**

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

The frozen verifier.

---

**`block_ids`** `torch.LongTensor`

Long tensor of shape \[batch, draft]; the block's token ids,
position 0 holding the verified anchor and the rest `MASK`.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel.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,
    kwargs = {}
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel.spec_generate(
    target: torch.nn.Module,
    input_ids: torch.LongTensor,
    max_new_tokens: int,
    stop_token_ids: list[int] | None,
    temperature: float,
    top_p: float = 1.0,
    top_k: int = 0
) -> torch.LongTensor
```

Block-parallel speculative decoding: draft a block, verify with the target, accept the matching prefix.

`top_p` / `top_k` truncate the *target's* distribution, which is what
defines the emitted tokens; the draft proposes from plain temperature.

**Parameters:**

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

The frozen verifier.

---

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

---

**`top_p`** `float` — default: 1.0

Nucleus mass to keep, in `(0, 1]`.

---

**`top_k`** `int` — default: 0

Candidates to keep, `0` for the whole vocabulary.

---

**Returns:** `torch.LongTensor`

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

```python
nemo_automodel.components.speculative.dflash.draft_qwen3._sliding_window_mask(
    query: torch.Tensor,
    key: torch.Tensor,
    sliding_window: int
) -> torch.Tensor
```

Band mask keeping keys within `sliding_window` of each query's position.

The draft's queries are the trailing `q_len` positions of the concatenated
`[context | noise-block]` key axis, so a query at row `i` sits at key
position `k_len - q_len + i`. Both bounds are strict, matching transformers'
`sliding_window_overlay` and the reference DFlash decode mask. The draft is
non-causal, so the band is symmetric; the forward half is inert in practice
because the block is far shorter than any real window.

**Parameters:**

**`query`** `torch.Tensor`

Tensor of shape \[batch, heads, query, head\_dim].

---

**`key`** `torch.Tensor`

Tensor of shape \[batch, kv\_heads, key, head\_dim], where `key`
spans the context followed by the draft block.

---

**`sliding_window`** `int`

Maximum absolute position distance, exclusive.

---

**Returns:** `torch.Tensor`

Floating-point tensor of shape \[1, 1, query, key], in `query.dtype`;

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.apply_rotary_pos_emb(
    q,
    k,
    cos,
    sin,
    unsqueeze_dim = 1
)
```

Apply RoPE where queries (draft block) are a suffix of the key positions.

The keys span `[context | noise-block]` while the queries are only the
noise block, so `q` is rotated with the trailing `q_len` slice of the
rotary tables and `k` with the full table.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.assert_target_supports_rollback(
    target: torch.nn.Module
) -> None
```

Reject targets whose cache cannot be rewound after a rejected block.

Speculative decoding verifies a whole block and then rewinds the target to
the accepted prefix. For attention layers that is just dropping KV rows, but
a linear-attention layer keeps a recurrent state that has already absorbed
the rejected tokens; `Cache.crop` truncates the KV entries and leaves that
state where it is. The target then predicts from a corrupted state, so the
"lossless" guarantee quietly stops holding -- greedy decoding drifts away
from the target's own output instead of reproducing it.

Training is unaffected: it is a single forward with no cache and no rewind.

**Parameters:**

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

The frozen verifier.

---

**Raises:**

* `ValueError`: If the target has any layer whose state cropping cannot
  rewind.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.build_qwen3_dflash_draft_config(
    target_config,
    num_draft_layers: int,
    num_target_layers: int,
    block_size: int,
    dflash_config: dict,
    attention_backend: str
) -> transformers.models.qwen3.configuration_qwen3.Qwen3Config
```

Build the DFlash draft config for a Qwen3-shaped target.

The draft is a small non-causal Qwen3 stack that reuses the target's
architecture defaults (head\_dim, rope\_theta, rms\_norm\_eps, ...) and only
shrinks the depth.

**Parameters:**

**`target_config`**

The target's config.

---

**`num_draft_layers`** `int`

Number of draft decoder layers.

---

**`num_target_layers`** `int`

Depth of the target; recorded so a reloaded draft
config still describes which target it was trained on.

---

**`block_size`** `int`

DFlash block size.

---

**`dflash_config`** `dict`

The recipe's DFlash block (`mask_token_id`,
`target_layer_ids`, and any subclass extras such as Domino's
projector fields).

---

**`attention_backend`** `str`

The draft's attention implementation; must agree with
the mask format the trainer builds.

---

**Returns:** `Qwen3Config`

The draft `Qwen3Config`.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.build_target_layer_ids(
    num_target_layers: int,
    num_draft_layers: int
) -> list[int]
```

Pick `num_draft_layers` target layers spread across the target's depth.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.extract_context_feature(
    hidden_states: list[torch.Tensor],
    layer_ids: list[int]
) -> torch.Tensor
```

Concatenate the selected target layers' hidden states along the feature dim.

`hidden_states` follows HF's `output_hidden_states` convention where
index 0 is the embedding output, so layer `i`'s output is at index
`i + 1`.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.resolve_output_head(
    target: torch.nn.Module
) -> torch.nn.Module
```

Return the target's output projection.

Prefers `lm_head` and falls back to `get_output_embeddings()`, which is
what a `*ForConditionalGeneration` target exposes instead.

**Parameters:**

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

The frozen verifier.

---

**Returns:** `nn.Module`

The module projecting hidden states to vocabulary logits.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.sample(
    logits: torch.Tensor,
    temperature: float = 0.0,
    top_p: float = 1.0,
    top_k: int = 0
) -> torch.Tensor
```

Greedy (`temperature == 0`) or truncated-distribution sampling over the last dim.

**Parameters:**

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

Tensor of shape \[..., vocab]; raw scores, with arbitrary leading
dimensions.

---

**`temperature`** `float` — default: 0.0

Sampling temperature; `0` selects greedily.

---

**`top_p`** `float` — default: 1.0

Nucleus mass to keep, in `(0, 1]`.

---

**`top_k`** `int` — default: 0

Candidates to keep, `0` for the whole vocabulary.

---

**Returns:** `torch.Tensor`

Long tensor of shape \[...]; the sampled token id per position.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.sampling_probs(
    logits: torch.Tensor,
    temperature: float,
    top_p: float = 1.0,
    top_k: int = 0
) -> torch.Tensor
```

Full-vocabulary sampling distribution after temperature, top-k, and top-p.

Speculative decoding needs the *distribution*, not just a draw: rejection
sampling compares the target's probability of a drafted token against the
draft's, so both sides must be built the same way. Filtered-out tokens keep a
probability of exactly 0 so the residual `clamp(p - q, 0)` stays correct.

**Parameters:**

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

Tensor of shape \[..., vocab]; raw scores, with arbitrary leading
dimensions.

---

**`temperature`** `float`

Sampling temperature; must be > 0.

---

**`top_p`** `float` — default: 1.0

Nucleus mass to keep, in `(0, 1]`.

---

**`top_k`** `int` — default: 0

Candidates to keep, `0` for the whole vocabulary.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., vocab] in float32; a probability distribution over

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.validate_sampling(
    temperature: float,
    top_p: float = 1.0,
    top_k: int = 0
) -> None
```

Reject sampling parameters the decode path cannot honor.

**Parameters:**

**`temperature`** `float`

Sampling temperature; `0` means greedy.

---

**`top_p`** `float` — default: 1.0

Nucleus mass to keep, in `(0, 1]`.

---

**`top_k`** `int` — default: 0

Candidates to keep, `0` for the whole vocabulary.

---

**Raises:**

* `ValueError`: If any parameter is outside its supported range.

```python
nemo_automodel.components.speculative.dflash.draft_qwen3.GREEDY_TEMPERATURE_EPS = 1e-05
```

```python
nemo_automodel.components.speculative.dflash.draft_qwen3._REWINDABLE_LAYER_TYPES = frozenset({'full_attention', 'sliding_attention', 'chunked_attention'})
```