> 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_glm_5_2

DSpark speculative-decoding draft model with a GLM-5.2 attention backbone.

This mirrors the DeepSeek V4 DSpark draft (`draft_deepseek_v4.py`) structurally;
only the attention internals and rotary plumbing are swapped for GLM-5.2's
Multi-head Latent Attention (MLA). The draft runs dense, non-causal attention over
a `[context | noise-block]` layout, with visibility supplied entirely by the
DFlash additive attention mask.

The GLM MLA is the classic DeepSeek-V3 layout (`q_a_proj`/`q_a_layernorm`/
`q_b_proj`, a compressed `kv_a_proj_with_mqa`/`kv_a_layernorm`/`kv_b_proj`
latent, a separate `o_proj`, and interleaved complex `freqs_cis` RoPE over the
`qk_rope_head_dim` slice), NOT V4's single shared K=V latent + grouped O-LoRA +
per-head sink. There is no DSA indexer / sparse-attention path here: the draft is
always dense, and GLM's sparse top-k machinery belongs to the target model, not
this draft.

## Module Contents

### Classes

| Name                                                                                                               | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`Glm5_2DSparkAttention`](#nemo_automodel-components-speculative-dspark-draft_glm_5_2-Glm5_2DSparkAttention)       | Dense, non-causal GLM MLA over a `[context \| noise-block]` layout.       |
| [`Glm5_2DSparkDecoderLayer`](#nemo_automodel-components-speculative-dspark-draft_glm_5_2-Glm5_2DSparkDecoderLayer) | Pre-norm residual block: GLM DSpark MLA followed by a dense SwiGLU MLP.   |
| [`Glm5_2DSparkMLP`](#nemo_automodel-components-speculative-dspark-draft_glm_5_2-Glm5_2DSparkMLP)                   | Dense SwiGLU MLP for the DSpark draft.                                    |
| [`Glm5_2DSparkModel`](#nemo_automodel-components-speculative-dspark-draft_glm_5_2-Glm5_2DSparkModel)               | DSpark draft model with a GLM-5.2 attention backbone (dense, non-causal). |

### Functions

| Name                                                                                                               | Description                                                           |
| ------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------- |
| [`_glm_5_2_eager_attention`](#nemo_automodel-components-speculative-dspark-draft_glm_5_2-_glm_5_2_eager_attention) | Dense eager attention over an additive mask (no GQA repeat, no sink). |

### Data

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

### API

```python
class nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkAttention(
    config,
    layer_idx: int
)
```

**Bases:** `Module`

Dense, non-causal GLM MLA over a `[context | noise-block]` layout.

Reuses GLM-5.2's DeepSeek-V3-style MLA projections (Q-LoRA, a compressed KV latent,
a separate value up-projection, interleaved complex RoPE on the `qk_rope_head_dim`
slice) but drops the DSA indexer (the draft is always dense). Queries come from the
noise block only, while the compressed KV latent spans context plus noise. Visibility
is set entirely by the DFlash additive `attention_mask` (no causal flag, no indexer).

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkAttention.forward(
    hidden_states: torch.Tensor,
    target_hidden_states: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> tuple[torch.Tensor, typing.Optional[torch.Tensor]]
```

```python
class nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkDecoderLayer(
    config,
    layer_idx: int
)
```

**Bases:** `Module`

Pre-norm residual block: GLM DSpark MLA followed by a dense SwiGLU MLP.

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkDecoderLayer.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,
    freqs_cis: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkMLP(
    config
)
```

**Bases:** `Module`

Dense SwiGLU MLP for the DSpark draft.

This intentionally avoids GLM's MoE and shared-expert paths; the draft keeps a
single dense feed-forward block sized to `moe_intermediate_size` (the per-expert
width), so the drafter stays lightweight relative to the 256-expert target.

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel(
    config
)
```

**Bases:** `Module`

DSpark draft model with a GLM-5.2 attention backbone (dense, non-causal).

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel._apply(
    fn,
    recurse = True
)
```

Keep the rotary `freqs` buffer in fp32 across dtype casts.

`model.to(bfloat16)` (the training build path) would otherwise round the RoPE
`freqs` table 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.

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel._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
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.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
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.initialize_embeddings_and_head(
    embed_tokens: torch.nn.Module,
    lm_head: torch.nn.Module,
    freeze: bool = True
)
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.predict_confidence_step(
    hidden_states: torch.Tensor,
    prev_token_ids: typing.Optional[torch.Tensor] = None
) -> typing.Optional[torch.Tensor]
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.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]
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.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]
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.Glm5_2DSparkModel.set_embedding_head_trainable(
    trainable: bool
)
```

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2._glm_5_2_eager_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: typing.Optional[torch.Tensor],
    scaling: float
) -> torch.Tensor
```

Dense eager attention over an additive mask (no GQA repeat, no sink).

GLM MLA keeps one K/V per query head (`num_key_value_heads == num_attention_heads`),
so there is no group repeat. The additive `attention_mask` uses `-inf` for masked
positions (the DFlash SDPA mask); a plain matmul + fp32 softmax handles it correctly
(the DFlash mask guarantees no fully-masked query row, so no row softmaxes to NaN).

**Parameters:**

`[B, H, q_len, qk_head_dim]`.

`[B, H, kv_len, qk_head_dim]`.

`[B, H, kv_len, v_head_dim]`.

`[B, 1, q_len, kv_len]` additive mask, or `None`.

softmax scale (`qk_head_dim ** -0.5`).

**Returns:** `torch.Tensor`

`[B, q_len, H, v_head_dim]` (heads folded back next to the feature dim).

```python
nemo_automodel.components.speculative.dspark.draft_glm_5_2.__all__ = ['Glm5_2DSparkModel', 'Glm5_2DSparkAttention', 'Glm5_2DSparkDecoderLayer']
```