nemo_automodel.components.models.deepseek_v41.dspark

Native DeepSeek V4.1 DSpark draft backbone.

The released draft is stored under mtp.0 through mtp.2 but is not the autoregressive MTP objective used by earlier DeepSeek models. It is a parallel five-position drafter whose three blocks use sliding-window MLA, routed MoE and single-pass mHC. This module implements the trainable, cache-free backbone; the shared frozen embedding/LM head and anchor sampling remain owned by the generic DSpark trainer.

Module Contents

Classes

Name	Description
`DeepseekV41DSparkBackbone`	Three-stage native drafter operating on prepared target and noise tensors.
`DeepseekV41DSparkBackboneOutput`	Native draft states consumed by the released output heads.
`DeepseekV41DSparkModel`	Model-owned adapter from the native draft to the shared DSpark loss contract.
`_DeepseekV41DSparkAttention`	Cache-free DSpark attention over target context and parallel draft blocks.
`_DeepseekV41DSparkBlock`	One released DSpark stage with MLA, MoE, mHC and stage-owned heads.
`_DeepseekV41DSparkConfidenceHead`	Predict conditional acceptance logits in FP32.
`_DeepseekV41DSparkMarkovHead`	Rank-factorized first-order token-transition bias.
`_DeepseekV41DSparkStageOutput`	Internal stage state kept within each stage’s FSDP forward boundary.

Data

__all__

API

class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None
)

Bases: Module

Three-stage native drafter operating on prepared target and noise tensors.

moe_config

mtp

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.build_attention_mask(
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    context_sequence: int,
    dtype: torch.dtype
) -> torch.Tensor

Build the released SWA-128 multi-anchor training mask.

Every query sees the target window ending immediately before its anchor, plus every parallel input in its own draft block. It cannot see another anchor’s block. Invalid padding blocks retain their own in-block keys so no attention row is fully masked; their losses are discarded later.

Parameters:

anchor_positions

torch.Tensor

Integer tensor [batch, num_anchors].

block_keep_mask

torch.Tensor

Boolean tensor [batch, num_anchors].

context_sequence

int

Number of target-context tokens.

dtype

torch.dtype

Floating dtype of the returned additive mask.

Returns: torch.Tensor

Additive tensor [batch, 1, num_anchors * block_size,

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.build_position_ids(
    anchor_positions: torch.Tensor,
    context_sequence: int
) -> torch.Tensor

Build official target-context and draft-query positions.

Parameters:

anchor_positions

torch.Tensor

Integer tensor [batch, num_anchors] containing the target token that seeds each draft block.

context_sequence

int

Number of target-context tokens.

Returns: torch.Tensor

Integer tensor [batch, context_sequence + num_anchors * block_size].

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.forward(
    noise_embeddings: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    previous_token_ids: torch.Tensor | None = None,
    enable_confidence_head: bool = True,
    confidence_head_stop_gradient: bool = False
) -> nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackboneOutput

Run the cache-free draft backbone for sampled anchors.

Parameters:

noise_embeddings

torch.Tensor

Tensor of shape [batch, draft_sequence, hidden], containing an anchor embedding followed by noise embeddings in each fixed-width block.

target_hidden_states

torch.Tensor

Concatenated target features of shape [batch, context_sequence, target_layers * hidden].

position_ids

torch.Tensor

Integer tensor of shape [batch, context_sequence + draft_sequence].

attention_mask

torch.Tensor

Additive tensor of shape [batch, 1, draft_sequence, context_sequence + draft_sequence].

previous_token_ids

torch.Tensor | NoneDefaults to None

Optional integer tensor of shape [batch, draft_sequence] used by the final Markov head.

enable_confidence_head

boolDefaults to True

Whether the final stage computes confidence.

confidence_head_stop_gradient

boolDefaults to False

Whether the confidence head reads detached inputs.

Returns: DeepseekV41DSparkBackboneOutput

Draft backbone output containing normalized states of shape [batch,

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackbone.initialize_weights(
    buffer_device: torch.device | None = None
) -> None

Initialize every draft parameter for checkpoint-free training.

Parameters:

buffer_device

torch.device | NoneDefaults to None

Device used by grouped expert initialization. Defaults to the first attention projection’s device.

class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkBackboneOutput(
    normalized_hidden_states: torch.Tensor,
    transition_logits: torch.Tensor | None = None,
    confidence_pred: torch.Tensor | None = None
)

Dataclass

Native draft states consumed by the released output heads.

confidence_pred

Tensor | None = None

normalized_hidden_states

Tensor

transition_logits

Tensor | None = None

class nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    num_anchors: int,
    enable_confidence_head: bool,
    confidence_head_stop_gradient: bool = False
)

Bases: DeepseekV41DSparkBackbone

Model-owned adapter from the native draft to the shared DSpark loss contract.

_keep_in_fp32_modules_strict

_no_split_modules

= ['_DeepseekV41DSparkBlock']

confidence_head_stop_gradient

= bool(confidence_head_stop_gradient)

embed_tokens

enable_confidence_head

= bool(enable_confidence_head)

layers

ModuleList

Expose native MTP stages to the shared AC/FSDP helpers.

lm_head

num_anchors

= int(num_anchors)

state_dict_adapter

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

Project hidden states through the frozen vocabulary head in FP32.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

Returns: torch.Tensor

FP32 tensor of shape [batch, sequence, vocab].

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.forward(
    input_ids: torch.Tensor,
    target_hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    target_last_hidden_states: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dspark.common.DSparkForwardOutput

Run native V4.1 DSpark training for sampled anchors.

Parameters:

input_ids

torch.Tensor

Target-token tensor of shape [batch, sequence].

target_hidden_states

torch.Tensor

Concatenated target-feature tensor of shape [batch, sequence, target_layers * hidden].

loss_mask

torch.Tensor

Supervision tensor of shape [batch, sequence].

target_last_hidden_states

torch.Tensor | NoneDefaults to None

Optional frozen target-state tensor of shape [batch, sequence, hidden] used by the probability-distance loss.

Returns: DSparkForwardOutput

Shared DSpark loss inputs. Draft logits have shape [batch,

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
) -> None

Copy the target’s embedding and vocabulary projection.

Parameters:

embed_tokens

nn.Module

Target embedding with weight of shape [vocab, hidden].

lm_head

nn.Module

Target output projection with weight of shape [vocab, hidden].

freeze

boolDefaults to True

Disable gradients for both copied modules when true.

nemo_automodel.components.models.deepseek_v41.dspark.DeepseekV41DSparkModel.set_embedding_head_trainable(
    trainable: bool
) -> None

Set whether the copied embedding and LM head receive gradients.

class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkAttention(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: DeepseekV41Attention

Cache-free DSpark attention over target context and parallel draft blocks.

nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor
) -> torch.Tensor

Attend draft queries to target context and their own parallel block.

Parameters:

hidden_states

torch.Tensor

Draft tensor of shape [batch, draft_sequence, hidden].

target_hidden_states

torch.Tensor

Projected target tensor of shape [batch, context_sequence, hidden].

position_ids

torch.Tensor

Integer tensor of shape [batch, context_sequence + draft_sequence], containing absolute positions for both regions.

attention_mask

torch.Tensor

Additive tensor broadcastable to shape [batch, heads, draft_sequence, context_sequence + draft_sequence], with zero for visible keys and negative infinity for masked keys.

Returns: torch.Tensor

Tensor of shape [batch, draft_sequence, hidden]. Inputs are not mutated.

class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkBlock(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    stage_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig
)

Bases: Module

One released DSpark stage with MLA, MoE, mHC and stage-owned heads.

attn

attn_hc

attn_norm

confidence_head

ffn

= MoE(moe_config, backend)

ffn_hc

ffn_norm

main_norm

main_proj

markov_head

mlp

MoE

Expose the shared parallelizer’s MoE interface without duplicate registration.

norm

nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkBlock.forward(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor,
    target_hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    previous_token_ids: torch.Tensor | None = None,
    enable_confidence_head: bool = True,
    confidence_head_stop_gradient: bool = False
) -> nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkStageOutput

Apply one native DSpark stage.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, draft_sequence, streams, hidden].

pre_mix

torch.Tensor

FP32 tensor of shape [batch, draft_sequence, streams].

target_hidden_states

torch.Tensor

Tensor of shape [batch, context_sequence, hidden].

position_ids

torch.Tensor

Integer tensor of shape [batch, context_sequence + draft_sequence].

attention_mask

torch.Tensor

Additive tensor of shape [batch, 1, draft_sequence, context_sequence + draft_sequence].

previous_token_ids

torch.Tensor | NoneDefaults to None

Optional integer tensor of shape [batch, draft_sequence] used by the final Markov head.

enable_confidence_head

boolDefaults to True

Whether the final stage computes confidence.

confidence_head_stop_gradient

boolDefaults to False

Whether the confidence head reads detached inputs, so its loss trains only confidence_head.

Returns: _DeepseekV41DSparkStageOutput

Stage state containing updated streams, target features, and any

class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkConfidenceHead(
    hidden_size: int,
    markov_rank: int
)

Bases: Module

Predict conditional acceptance logits in FP32.

AutoModel-trained drafts feed the final RMSNorm output to this head. Serving these checkpoints must use the same input instead of the released raw residual.

proj

nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkConfidenceHead.forward(
    hidden_states: torch.Tensor,
    markov_embeddings: torch.Tensor
) -> torch.Tensor

Predict a conditional acceptance logit for every draft position.

Parameters:

hidden_states

torch.Tensor

Final RMSNorm output of shape […, hidden].

markov_embeddings

torch.Tensor

Tensor of shape […, markov_rank] with matching leading dimensions.

Returns: torch.Tensor

FP32 tensor of shape […] containing uncalibrated confidence logits.

class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkMarkovHead(
    vocab_size: int,
    rank: int,
    dtype: torch.dtype
)

Bases: Module

Rank-factorized first-order token-transition bias.

embed

= nn.Embedding(vocab_size, rank, dtype=dtype)

head

nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkMarkovHead.forward(
    token_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Compute the transition bias and its conditioning embedding.

Parameters:

token_ids

torch.Tensor

Integer tensor of shape […] containing preceding tokens.

Returns: torch.Tensor

Transition logits of shape […, vocab] and embeddings of shape

class nemo_automodel.components.models.deepseek_v41.dspark._DeepseekV41DSparkStageOutput()

Bases: NamedTuple

Internal stage state kept within each stage’s FSDP forward boundary.

confidence_pred

Tensor | None = None

normalized_hidden_states

Tensor | None = None

pre_mix

Tensor

streams

Tensor

target_hidden_states

Tensor

transition_logits

Tensor | None = None

nemo_automodel.components.models.deepseek_v41.dspark.__all__ = ['DeepseekV41DSparkModel']