nemo_automodel.components.speculative.dflash.draft_qwen3

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.

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]

Bases: GradientCheckpointingLayer

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

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

Bases: Qwen3PreTrainedModel

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

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

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

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

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,
    top_p: float = 1.0,
    top_k: int = 0
) -> torch.LongTensor

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:

Returns: torch.Tensor

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

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.

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:

Raises:

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

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:

Returns: Qwen3Config

The draft Qwen3Config.

Pick num_draft_layers target layers spread across the target’s depth.

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.

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:

Returns: nn.Module

The module projecting hidden states to vocabulary logits.

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

Parameters:

Returns: torch.Tensor

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

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:

Returns: torch.Tensor

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

Reject sampling parameters the decode path cannot honor.

Parameters:

Raises:

• ValueError: If any parameter is outside its supported range.