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`	Non-causal attention whose keys/values are `[context \| noise-block]`.
`Qwen3DFlashDecoderLayer`	A DFlash decoder block: non-causal attention over `[context \| noise]` + MLP.
`Qwen3DFlashDraftModel`	DFlash draft model: a small non-causal Qwen3 stack over `[context \| noise]`.

Functions

Name	Description
`_sliding_window_mask`	Band mask keeping keys within `sliding_window` of each query’s position.
`apply_rotary_pos_emb`	Apply RoPE where queries (draft block) are a suffix of the key positions.
`assert_target_supports_rollback`	Reject targets whose cache cannot be rewound after a rejected block.
`build_target_layer_ids`	Pick `num_draft_layers` target layers spread across the target’s depth.
`extract_context_feature`	Concatenate the selected target layers’ hidden states along the feature dim.
`sample`	Greedy (temperature ~ 0) or temperature sampling over the last dim.

Data

_REWINDABLE_LAYER_TYPES

API

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

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]

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

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

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

gru_hidden_dim

= dflash_config['gru_hidden_dim']

hidden_norm

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

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.

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

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

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

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;

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.

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.

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.

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.

nemo_automodel.components.speculative.dflash.draft_qwen3.sample(
    logits: torch.Tensor,
    temperature: float = 0.0
) -> torch.Tensor

Greedy (temperature ~ 0) or temperature sampling over the last dim.

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