Module deepseek_v4_mtp

Module deepseek_v4_mtp 

Source
Expand description

DeepSeek-V4-Flash Multi-Token-Prediction (MTP) draft proposer.

Implements DraftProposer over the DeepseekV4MtpModule loaded by load_v4_mtp_module. Unlike the Qwen-shaped crate::layers::MtpHead (a hand-rolled single attention + MoE block), the V4 MTP module’s body is a full reused V4 layer (MLA + manifold-constrained hyper-connections (mHC) + 256-expert NVFP4 MoE). The proposer therefore delegates the bulk of the forward to body.decode() and only wraps it with the MTP-specific pieces.

Forward (propose(), K = 1 since num_nextn_predict_layers == 1):

  embed   = embed_tokens[last_token]                       // [hidden] BF16
  h_in    = e_proj · rms_norm(embed,  enorm)
          + h_proj · rms_norm(hidden, hnorm)               // combiner
  hc_expand(h_in → hc_streams)                             // is_first mHC
  body.decode(hc_streams, …, mtp_kv_cache, state.seq_len)  // MIDDLE mHC + MLA + MoE
  hc_head(hc_streams → h_out)                              // is_last mHC
  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 sees is_first_layer == false AND is_last_layer == false: it runs the middle mHC mixing (hc_pre → attn → hc_post → hc_pre → ffn → hc_post) reading/writing hc_streams, but does NOT call hc_expand or hc_head. The proposer supplies both ends.

§Separate KV cache + distinct metadata offset

The MTP attention writes into its OWN single-layer MLA-shaped PagedKvCache (num_kv_heads = 1, head_dim = kv_lora_rank + qk_rope_head_dim), never the target’s. The V4-Flash decode attention (attention_forward_v4) reads positions / slot / seq_len / block_table from ctx.attn_metadata, so the proposer uploads MTP-specific metadata to scratch().offset(MTP_META_OFFSET) — distinct from the target metadata at 32768 — and threads it through a derived ForwardContext.

Structs§

DeepseekV4MtpHead
DeepSeek-V4 MTP draft proposer.
DeepseekV4MtpProposerState
Per-sequence state for the DeepSeek-V4 MTP proposer.