Module glm5next_mtp_head 

Glm5NextMtpHead — GLM-5.3’s MTP block as a DraftProposer.

One draft token per forward_one:

x = eh_proj( concat( enorm(embed[token]), hnorm(target_hidden) ) )   [1, 2H] -> [1, H]
x = layers.45(x)                       DSA + routed MoE, PLAIN residual (no mHC)
logits = lm_head( shared_head.norm(x) )              the target's own BF16 head
draft  = argmax(logits)

🔴 This block is SHARDED — EP-sharded routed MoE (144 of 288 experts per rank) and a row-parallel DSA o_proj — unlike the Qwen and DeepSeek-V4 MTP modules, which load every expert on every rank. So it needs the communicator exactly as a text layer does.

🪤 Historically it ran WITHOUT one and on RANK 0 ONLY (run_mtp_propose_multi_dispatch: “Rank 1 does not participate in MTP propose”), which is correct for V4 and wrong here: the drafter proposed from half the routed sum and half the attention output. Lossless — the target verifies every draft — so the only symptom was acceptance. ATLAS_MTP_EP_PROPOSE=1 turns on BOTH halves of the fix: the worker executes propose on EP_CMD_MTP_PROPOSE, and needs_comm() then hands the block a comm. Turning on only the second half is t58, which deadlocked.

🪤 The embedding is read as a POINTER into the shared table, not a gather: the row for token t is embed_tokens + t * hidden * 2. No kernel, no copy.

Structs

Glm5NextMtpHead

Glm5NextMtpProposerState

Per-sequence drafter state: the block’s own indexer cache and KV blocks.