AI Engineer Exhaustive
MediumPlusGLM-4.5

MTP Head

GLM-4.5

Medium

Multi-Token Prediction head used in GLM-4.5 to predict additional future tokens for speculative decoding.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

GLM-4.5 attaches one extra MoE transformer layer as a multi-token prediction (MTP) head. Table 1 lists # MTP Layers = 1 for both GLM-4.5 and Air. §2.1 says the layer exists “to support speculative decoding during inference,” citing the Gloeckle et al. MTP objective. §2.4 gives the only training hyperparameter: the MTP loss weight (\lambda) is (0.3) for the first 15T tokens and (0.1) afterwards.

The paper does not write the MTP residual graph. That graph is specified by the official inference code the README points at — vLLM glm4_moe_mtp.py and SGLang’s next-n path — and it matches DeepSeek-V3’s MTP module. Hugging Face Transformers explicitly does not implement this head; causal-LM logits there come only from the main stack.

So: the paper fixes that there is one MoE MTP layer and why (speculative decoding + an auxiliary loss). The forward equations below are inferred from those official kernels, not from a numbered equation in the PDF.

How it works

The main model produces hidden states (h_t \in \mathbb{R}^{D}) at each position after the final RMSNorm of the (L)-layer stack ((L=92) or (46)). MTP’s job at training time is to predict token (t+2) given (h_t) and the embedding of the already-known next token (e_{t+1}). At decode time the same module drafts extra future tokens for speculative verification.

Let (e_{t+1} = E[t_{t+1}]) share the backbone embedding table. Two independent RMSNorms and a fusion projection produce the MTP residual:

u_t = W_{eh}\,\big[\,\mathrm{RMSNorm}_e(e_{t+1});\,\mathrm{RMSNorm}_h(h_t)\,\big], \qquad W_{eh} \in \mathbb{R}^{D \times 2D}.

(u_t) then goes through one full GLM-4.5 decoder block — the same pre-norm attention + MoE FFN as a sparse backbone layer, including group-routed top-8 and the shared expert:

z_t = \mathrm{Block}_{\mathrm{MTP}}(u_t).

A small “shared head” applies another RMSNorm and the output vocabulary projection (loaded from / shared with lm_head in the vLLM weight map):

\hat{y}_{t+2} = W_{\mathrm{lm}}\,\mathrm{RMSNorm}(z_t).

Checkpoints store this module as model.layers.{L} (num_nextn_predict_layers = 1), with extra tensors enorm, hnorm, eh_proj, shared_head, plus a full mtp_block (attention, norms, MoE). vLLM rewrites those names so the transformer submodules sit under .mtp_block. and embeddings are shared with the main model.

During speculative decoding, vLLM / SGLang call this head with spec_step_idx and the previous hidden state; the drafted token is verified against the backbone in the usual accept/reject loop. The README’s vLLM launch line uses --speculative-config.method mtp.

What is specified vs inferred:

Specified in the paper Inferred from official inference code
1 MTP layer, MoE, both model sizes enorm / hnorm / eh_proj fusion
Used for speculative decoding One full decoder block as mtp_block
Loss weight (\lambda \in {0.3, 0.1}) Shared embedding + shared / remapped lm_head
Counted in the 355B / 106B totals Layer index (L) in the checkpoint

The exact training loss (whether (\lambda) multiplies a next-next cross-entropy only, and whether more than one future step is trained) is not written. The public kernels implement a single extra token.

Official code

The zai-org README states that MTP lives in vLLM glm4_moe_mtp.py and SGLang glm4_moe.py (Glm4MoeForCausalLMNextN / next-n weight remap). The vLLM types to read are Glm4MoeMultiTokenPredictorLayer (enorm, hnorm, eh_proj, mtp_block, shared_head) and Glm4MoeMultiTokenPredictor.compute_logits.

Transformers Glm4MoeForCausalLM has no MTP module; its docs say so. zai-org/GLM-4.5 has no modeling source of its own.

Watch-outs

Sources