nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2

DFlash 2 draft model (Qwen3-style).

DFlash 2 (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`	Pairwise path selector over each draft position’s top-k candidates.
`GroupedDynamicCausalConv`	Two-tap dynamic depthwise convolution wrapped around one draft sublayer.
`Qwen3DFlash2DecoderLayer`	A DFlash 2 decoder block: DFlash’s block + a two-tap conv around each sublayer.
`Qwen3DFlash2DraftModel`	DFlash 2 draft model: the DFlash stack plus in-block convs and a path selector.

Functions

Name	Description
`_grouped_dynamic_convolve`	Depthwise convolution over draft-block positions with content-adaptive taps.
`dflash2_rejection_sample`	Accept a prefix of the drafted block and resample the first rejected token.

API

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) + <A(a) * H(h_t), B(b)> 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))

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

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.

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

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

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

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

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

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

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

Bases: Qwen3DFlashDecoderLayer

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

attention_conv

mlp_conv

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 | NoneDefaults to None

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

hidden_states

torch.Tensor | NoneDefaults to 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 | NoneDefaults to None

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

position_ids

torch.LongTensor | NoneDefaults to None

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

past_key_value

Cache | NoneDefaults to None

Draft KV cache, or None.

use_cache

bool | NoneDefaults to False

Whether to write into past_key_value.

cache_position

torch.LongTensor | NoneDefaults to None

Long tensor of shape [draft], or None.

position_embeddings

tuple[torch.Tensor, torch.Tensor] | NoneDefaults to None

Tuple of rotary (cos, sin) tensors of shape [batch, context + draft, head_dim].

conv_block_size

intDefaults to 1

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

**kwargs

Defaults to {}

Forwarded to the attention implementation.

Returns: torch.Tensor

Tensor of shape [batch, draft, hidden].

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

Bases: Qwen3DFlashDraftModel

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

_no_split_modules

= ['Qwen3DFlash2DecoderLayer']

candidate_selector

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 | NoneDefaults to None

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

noise_embedding

torch.Tensor | NoneDefaults to None

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

target_hidden

torch.Tensor | NoneDefaults to None

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

past_key_values

Cache | NoneDefaults to None

Draft KV cache, or None.

use_cache

boolDefaults to False

Whether to write into past_key_values.

conv_block_size

int | NoneDefaults to None

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

**kwargs

Defaults to {}

Forwarded to the attention implementation.

Returns: torch.Tensor

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

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.

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

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

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