> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.dspark.draft_kimi_k3

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

## Module Contents

### Classes

| Name                                                                                                               | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`KimiK3DSparkAttention`](#nemo_automodel-components-speculative-dspark-draft_kimi_k3-KimiK3DSparkAttention)       | Dense non-causal K3 MLA over target context and a parallel noise block. |
| [`KimiK3DSparkDecoderLayer`](#nemo_automodel-components-speculative-dspark-draft_kimi_k3-KimiK3DSparkDecoderLayer) | Pre-norm K3 MLA and dense SiTU feed-forward block.                      |
| [`KimiK3DSparkModel`](#nemo_automodel-components-speculative-dspark-draft_kimi_k3-KimiK3DSparkModel)               | DSpark draft with a dense K3 MLA backbone and no KDA or MoE layers.     |

### Data

[`__all__`](#nemo_automodel-components-speculative-dspark-draft_kimi_k3-__all__)

### API

```python
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)`

---

```python
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]` — default: None

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

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, query_length, hidden]`.

```python
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)`

---

```python
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]` — default: None

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

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, query_length, hidden]`.

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

**Bases:** [DeepseekV4DSparkModel](/nemo-automodel/nemo_automodel/components/speculative/dspark/draft_deepseek_v4#nemo_automodel-components-speculative-dspark-draft_deepseek_v4-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)`

---

**`fc`**

---

**`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`

---

```python
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]` — default: None

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

---

**`noise_embedding`** `Optional[torch.Tensor]` — default: None

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

---

**`target_hidden_states`** `Optional[torch.Tensor]` — default: 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]`.

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