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 SomeThe 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§
- Deepseek
V4Mtp Head - DeepSeek-V4 MTP draft proposer.
- Deepseek
V4Mtp Proposer State - Per-sequence state for the DeepSeek-V4 MTP proposer.