nemo_automodel.components.speculative.dspark.draft_deepseek_v4

DSpark speculative-decoding draft model with a DeepSeek V4 attention backbone.

This mirrors the Qwen3 DSpark draft (draft_qwen3.py) structurally; only the attention internals, MLP, and rotary plumbing are swapped for DeepSeek V4. The draft runs dense, non-causal attention over a [context | noise-block] layout, with visibility supplied entirely by the DFlash additive attention mask. There is no compressor / indexer / sparse-attention path (the draft is always dense; V4’s sparse machinery belongs to the target model, not this draft).

Module Contents

Classes

Name	Description
`DeepseekV4DSparkAttention`	Dense, non-causal V4 attention over a `[context \| noise-block]` layout.
`DeepseekV4DSparkDecoderLayer`	Pre-norm residual block: V4 DSpark attention followed by a dense SwiGLU MLP.
`DeepseekV4DSparkMLP`	Dense SwiGLU MLP for the DSpark draft.
`DeepseekV4DSparkModel`	DSpark draft model with a DeepSeek V4 attention backbone (dense, non-causal).

Data

__all__

API

class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkAttention(
    config,
    layer_idx: int
)

Bases: Module

Dense, non-causal V4 attention over a [context | noise-block] layout.

Reuses DeepSeek V4’s Q-LoRA, single shared K=V latent, grouped O-LoRA, and the per-head attention sink. Queries come from the noise block only, while the shared latent spans context plus noise. Visibility is set entirely by the DFlash additive attention_mask (there is no causal flag, and no compressor or indexer, since the draft is always dense).

attention_dropout

head_dim

= int(config.head_dim)

kv_norm

layer_idx

= int(layer_idx)

num_heads

= int(config.num_attention_heads)

num_key_value_groups

= int(config.num_attention_heads)

q_norm

rope_head_dim

= int(config.qk_rope_head_dim)

scaling

= self.head_dim ** -0.5

sinks_param

wkv

wo_a

wo_b

wq_a

wq_b

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> tuple[torch.Tensor, typing.Optional[torch.Tensor]]

class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkDecoderLayer(
    config,
    layer_idx: int
)

Bases: Module

Pre-norm residual block: V4 DSpark attention followed by a dense SwiGLU MLP.

hidden_size

= config.hidden_size

input_layernorm

mlp

= DeepseekV4DSparkMLP(config)

post_attention_layernorm

self_attn

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkDecoderLayer.forward(
    target_hidden_states: typing.Optional[torch.Tensor] = None,
    hidden_states: typing.Optional[torch.Tensor] = None,
    attention_mask: typing.Optional[torch.Tensor] = None,
    position_ids: typing.Optional[torch.LongTensor] = None,
    past_key_value: typing.Optional[object] = None,
    use_cache: typing.Optional[bool] = False,
    position_embeddings: typing.Optional[tuple[torch.Tensor, torch.Tensor]] = None,
    kwargs = {}
) -> torch.Tensor

class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkMLP(
    config
)

Bases: Module

Dense SwiGLU MLP for the DSpark draft.

This intentionally avoids V4’s MoE and Hyper-Connection paths; the draft keeps a plain 3D hidden state and a single dense feed-forward block.

act_fn

= ACT2FN[config.hidden_act]

down_proj

gate_proj

up_proj

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel(
    config
)

Bases: Module

DSpark draft model with a DeepSeek V4 attention backbone (dense, non-causal).

_no_split_modules

= ['DeepseekV4DSparkDecoderLayer']

block_size

= int(config.block_size)

embed_tokens

enable_confidence_head

= bool(config.enable_confidence_head)

hidden_norm

layers

lm_head

markov_head

= build_markov_head(config)

mask_token_id

= config.mask_token_id

norm

num_anchors

= int(config.num_anchors)

partial_rotary_factor

rotary_emb

target_layer_ids

= config.target_layer_ids

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel._apply(
    fn,
    recurse = True
)

Keep the rotary inv_freq buffer in fp32 across dtype casts.

model.to(bfloat16) (the training build path) would otherwise round inv_freq to bf16 and dephase RoPE with absolute position, eroding draft acceptance (the mismatch grows with position, and a bf16 round-trip cannot be undone by upcasting). Snapshot the fp32 frequencies before the cast and restore them after (the Fp32Safe rotary idiom used elsewhere in the repo), so the buffer never makes a bf16 round-trip.

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel._forward_backbone(
    position_ids: torch.LongTensor,
    attention_mask: typing.Optional[torch.Tensor] = None,
    noise_embedding: typing.Optional[torch.Tensor] = None,
    target_hidden_states: typing.Optional[torch.Tensor] = None,
    past_key_values: typing.Optional[object] = None,
    use_cache: bool = False,
    kwargs = {}
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_last_hidden_states: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
)

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.predict_confidence_step(
    hidden_states: torch.Tensor,
    prev_token_ids: typing.Optional[torch.Tensor] = None
) -> typing.Optional[torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.sample_draft_token_step(
    base_logits: torch.Tensor,
    prev_token_ids: torch.Tensor,
    temperature: float = 0.0,
    hidden_states: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.sample_draft_tokens(
    base_logits: torch.Tensor,
    first_prev_token_ids: torch.Tensor,
    temperature: float = 0.0,
    hidden_states: typing.Optional[torch.Tensor] = None
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.DeepseekV4DSparkModel.set_embedding_head_trainable(
    trainable: bool
)

nemo_automodel.components.speculative.dspark.draft_deepseek_v4.__all__ = ['DeepseekV4DSparkModel', 'DeepseekV4DSparkAttention', 'DeepseekV4DSparkDecoderL...