Module deepseek_v4_mtp

Module deepseek_v4_mtp 

Source
Expand description

2026-09-25: DeepSeek-V4-Flash multi-token-prediction (MTP) draft proposer.

Implements DraftProposer over the DeepseekV4MtpModule loaded by load_v4_mtp_module. Where metrale_model_layers::layers::MtpHead builds its own attention + MoE block, the V4 MTP module’s body is a reused V4 layer (MLA + mHC hyper-connections + MoE), so the proposer delegates the bulk of the forward to body.decode() and wraps it with the MTP-specific pieces.

One draft (forward_one; propose runs it num_drafts times):

  embed   = embed_tokens[token]                            // [hidden] bf16
  h_in    = e_proj · rms_norm(embed,  enorm)
          + h_proj · rms_norm(hidden, hnorm)               // combiner
  hc_expand(h_in → hc_streams)                             // first mHC stage
  body.decode(hc_streams, …, mtp_kv_cache, state.seq_len)  // middle mHC + MLA + MoE
  hc_head(hc_streams → h_out)                              // last mHC stage
  logits  = lm_head(rms_norm(h_out, norm))
  draft   = argmax(logits)                                 // grammar-masked when Some

The body was assembled with layer_idx = num_hidden_layers, so its decode_inner_hc is neither the first nor the last model layer: it runs the middle mHC mixing (hc_pre → attn → hc_post → hc_pre → ffn → hc_post) on hc_streams and calls neither hc_expand nor hc_head. The proposer supplies both ends.

§Separate KV cache and metadata

The MTP attention writes into its own MLA-shaped PagedKvCache (num_kv_heads = 1, head_dim = kv_lora_rank + qk_rope_head_dim), never the target’s. The decode attention reads positions / slot / seq_len / block_table from ctx.attn_metadata, so the proposer uploads its own metadata to scratch().offset(MTP_META_OFFSET), away from the target’s at 32768, and passes it in a derived ForwardContext.

Owner: model-arch, DeepSeek-V4 MTP. Invariants: none beyond the types.

Structs§

DeepseekV4MtpHead
2026-09-25: DeepSeek-V4 MTP draft proposer. embed_tokens and lm_head are the target model’s bf16 tables; mtp_vocab_size limits the draft LM-head GEMV (0 = full vocab); kv_cache is the MTP attention’s own cache.
DeepseekV4MtpProposerState
2026-09-25: Per-sequence state for the DeepSeek-V4 MTP proposer: the block table and length of the sequence in the MTP KV cache, the draft count of the last propose (for after_verify), and the body’s layer state from body.alloc_state.