nemo_automodel.components.speculative.dspark.draft_kimi_k3

Dense Kimi K3 MLA backbone for DSpark speculative-decoding training.

Module Contents

Classes

Name	Description
`KimiK3DSparkAttention`	Dense non-causal K3 MLA over target context and a parallel noise block.
`KimiK3DSparkDecoderLayer`	Pre-norm K3 MLA and dense SiTU feed-forward block.
`KimiK3DSparkModel`	DSpark draft with a dense K3 MLA backbone and no KDA or MoE layers.

Data

__all__

API

class nemo_automodel.components.speculative.dspark.draft_kimi_k3.KimiK3DSparkAttention(
    config,
    layer_idx: int
)

Bases: Module

Dense non-causal K3 MLA over target context and a parallel noise block.

K3 MLA keeps one K/V per query head (num_key_value_heads == num_attention_heads), like GLM-5.2’s MLA, so there is no GQA group repeat: kv_b_proj already emits one qk_nope_head_dim + v_head_dim slice per query head.

attention_dropout

g_proj

kv_a_layernorm

kv_a_proj_with_mqa

kv_b_proj

kv_lora_rank

= int(config.kv_lora_rank)

layer_idx

= int(layer_idx)

num_heads

= int(config.num_attention_heads)

num_key_value_heads

= int(config.num_key_value_heads)

o_proj

q_a_layernorm

q_a_proj

q_b_proj

q_head_dim

= self.qk_nope_head_dim + self.qk_rope_head_dim

q_lora_rank

= config.q_lora_rank

q_proj

qk_nope_head_dim

= int(config.qk_nope_head_dim)

qk_rope_head_dim

= int(config.qk_rope_head_dim)

scaling

= self.q_head_dim ** -0.5

use_output_gate

= bool(config.mla_use_output_gate)

v_head_dim

= int(config.v_head_dim)

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

Attend from the noise block to the frozen target context plus itself.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, query_length, hidden]; the noise block that supplies the queries.

target_hidden_states

torch.Tensor

Tensor of shape [batch, context_length, hidden]; the frozen target context, concatenated before hidden_states to form the keys and values.

attention_mask

Optional[torch.Tensor]Defaults to None

Optional additive mask of shape [batch, 1, query_length, context_length + query_length].

Returns: torch.Tensor

Tensor of shape [batch, query_length, hidden].

class nemo_automodel.components.speculative.dspark.draft_kimi_k3.KimiK3DSparkDecoderLayer(
    config,
    layer_idx: int
)

Bases: Module

Pre-norm K3 MLA and dense SiTU feed-forward block.

input_layernorm

mlp

= KimiK3MLP(config, dtype=(torch.float32))

post_attention_layernorm

self_attn

= KimiK3DSparkAttention(config, layer_idx)

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

Run one draft layer over the noise block against the target context.

Parameters:

target_hidden_states

torch.Tensor

Tensor of shape [batch, context_length, hidden]; the frozen target context supplying the extra keys and values.

hidden_states

torch.Tensor

Tensor of shape [batch, query_length, hidden]; the noise block, also the residual stream.

attention_mask

Optional[torch.Tensor]Defaults to None

Optional additive mask of shape [batch, 1, query_length, context_length + query_length].

Returns: torch.Tensor

Tensor of shape [batch, query_length, hidden].

class nemo_automodel.components.speculative.dspark.draft_kimi_k3.KimiK3DSparkModel(
    config
)

Bases: DeepseekV4DSparkModel

DSpark draft with a dense K3 MLA backbone and no KDA or MoE layers.

_no_split_modules

= ['KimiK3DSparkDecoderLayer']

block_size

= int(config.block_size)

confidence_head_with_markov

embed_tokens

enable_confidence_head

= bool(config.enable_confidence_head)

hidden_norm

layers

lm_head

markov_head

= build_markov_head(config)

mask_token_id

= int(config.mask_token_id)

norm

num_anchors

= int(config.num_anchors)

target_layer_ids

= config.target_layer_ids

nemo_automodel.components.speculative.dspark.draft_kimi_k3.KimiK3DSparkModel._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,
    kwargs = {}
) -> torch.Tensor

Project the target context once, then run the draft layers over the noise block.

Parameters:

position_ids

torch.LongTensor

Tensor of shape [batch, sequence]; unused (NoPE MLA).

attention_mask

Optional[torch.Tensor]Defaults to None

Optional additive mask of shape [batch, 1, query_length, context_length + query_length].

noise_embedding

Optional[torch.Tensor]Defaults to None

Tensor of shape [batch, query_length, hidden]; the noise block that becomes the queries.

target_hidden_states

Optional[torch.Tensor]Defaults to None

Tensor of shape [batch, context_length, len(target_layer_ids) * hidden]; the concatenated target hidden states, projected by fc and normed to [batch, context_length, hidden].

Returns: torch.Tensor

Tensor of shape [batch, query_length, hidden].

nemo_automodel.components.speculative.dspark.draft_kimi_k3.__all__ = ['KimiK3DSparkAttention', 'KimiK3DSparkDecoderLayer', 'KimiK3DSparkModel']