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`	Dense, non-causal GLM MLA over a `[context \| noise-block]` layout.
`Glm5_2DSparkDecoderLayer`	Pre-norm residual block: GLM DSpark MLA followed by a dense SwiGLU MLP.
`Glm5_2DSparkMLP`	Dense SwiGLU MLP for the DSpark draft.
`Glm5_2DSparkModel`	DSpark draft model with a GLM-5.2 attention backbone (dense, non-causal).

Functions

Name	Description
`_glm_5_2_eager_attention`	Dense eager attention over an additive mask (no GQA repeat, no sink).

Data

__all__

API

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).

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)

o_proj

q_a_layernorm

q_a_proj

q_b_proj

q_lora_rank

= int(config.q_lora_rank)

qk_head_dim

= self.qk_nope_head_dim + self.qk_rope_head_dim

qk_nope_head_dim

= int(config.qk_nope_head_dim)

qk_rope_head_dim

= int(config.qk_rope_head_dim)

scaling

= self.qk_head_dim ** -0.5

v_head_dim

= int(config.v_head_dim)

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

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.

hidden_size

= config.hidden_size

input_layernorm

mlp

= Glm5_2DSparkMLP(config)

post_attention_layernorm

self_attn

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

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.

act_fn

= ACT2FN[config.hidden_act]

down_proj

gate_proj

up_proj

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

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).

_no_split_modules

= ['Glm5_2DSparkDecoderLayer']

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)

target_layer_ids

= config.target_layer_ids

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.

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

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

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

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
)

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]

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]

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]

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

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:

query

torch.Tensor

[B, H, q_len, qk_head_dim].

key

torch.Tensor

[B, H, kv_len, qk_head_dim].

value

torch.Tensor

[B, H, kv_len, v_head_dim].

attention_mask

Optional[torch.Tensor]

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

scaling

float

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).

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