metrale_model_arch/glm5next_mtp_head.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Glm5NextMtpHead`: GLM-5.3's MTP block as a [`DraftProposer`], one draft token per `forward_one`.
4//!
5//! Owner: model-arch (GLM-5.3 MTP drafter).
6//! Invariants:
7//! - The FP8 copy of the draft head is read only on the vocab-sharded sweep; every other
8//! sweep reads the BF16 `lm_head`.
9//!
10//! ```text
11//! x = eh_proj( concat( enorm(embed[token]), hnorm(target_hidden) ) ) [1, 2H] -> [1, H]
12//! x = layers.{num_hidden_layers}(x) DSA + routed MoE, plain residual (no mHC)
13//! logits = lm_head( shared_head.norm(x) ) the target's own head
14//! draft = argmax(logits)
15//! ```
16//!
17//! The block is sharded like a text layer: its routed experts are split across EP ranks and
18//! its DSA `o_proj` is row-parallel, so its forward needs the communicator (`needs_comm`).
19
20use anyhow::{Result, bail};
21use parking_lot::Mutex;
22use std::any::Any;
23
24use metrale_cache::kv_cache::{KvCacheConfig, KvCacheDtype, PagedKvCache};
25use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
26use metrale_gpu_runtime::kernel_args::KernelLaunch;
27
28use crate::glm5next_dsa::state::Glm5NextDsaState;
29use metrale_model_layers::layer::{ForwardContext, LayerState};
30
31use crate::weight_loader::Glm5NextMtpModule;
32use metrale_model_layers::layers::ops;
33use metrale_model_layers::speculative::{DraftProposer, ProposerState};
34use metrale_model_layers::weight_map::DenseWeight;
35
36mod init;
37mod proposer;
38
39/// 2026-09-25: Per-sequence drafter state: the block's own indexer cache and KV block table.
40pub struct Glm5NextMtpProposerState {
41 dsa: Glm5NextDsaState,
42 /// 2026-09-25: Rows the drafter has written. `propose` and `after_verify` roll it back.
43 seq_len: usize,
44 block_table: Vec<u32>,
45 /// 2026-09-25: How many drafts the last `propose` wrote, so `after_verify` knows what
46 /// to trim.
47 last_drafted: usize,
48 /// 2026-09-25: Scratch: `[2, hidden]` BF16 concat, `[hidden]` BF16 block input,
49 /// `[vocab]` BF16 logits, `[1]` u32 argmax.
50 concat: DevicePtr,
51 x: DevicePtr,
52 logits: DevicePtr,
53 arg: DevicePtr,
54 /// 2026-09-25: 8 BF16 lanes for the vocab-sharded head's cross-rank pick: each rank's
55 /// max logit, then three base-256 digits of each rank's token id.
56 head_xchg: DevicePtr,
57 /// 2026-09-25: Set by `free_state`, which is a no-op once it is set.
58 released: bool,
59}
60
61impl ProposerState for Glm5NextMtpProposerState {
62 fn as_any(&self) -> &dyn Any {
63 self
64 }
65 fn as_any_mut(&mut self) -> &mut dyn Any {
66 self
67 }
68}
69
70pub struct Glm5NextMtpHead {
71 module: Glm5NextMtpModule,
72 embed_tokens: DenseWeight,
73 lm_head: DenseWeight,
74 /// 2026-09-25: A one-layer pool of the drafter's own, so its entries can be trimmed
75 /// independently of the target's.
76 kv_cache: Mutex<PagedKvCache>,
77 rms_norm_k: KernelHandle,
78 gemv_k: KernelHandle,
79 argmax_k: KernelHandle,
80 hidden: usize,
81 vocab: usize,
82 max_seq_len: usize,
83 /// 2026-09-25: Vocab shard of the shared `lm_head` this rank sweeps:
84 /// `[head_v0, head_v0 + head_n)`. `head_n == vocab` on a single rank or when `vocab_size`
85 /// does not divide by `tp_world_size`.
86 head_rank: usize,
87 head_v0: usize,
88 head_n: usize,
89 /// 2026-09-25: FP8 E4M3 copy of this rank's vocab shard of `lm_head`, for drafting only.
90 /// The target verifies every draft with its own head, so this copy can change the
91 /// acceptance rate but not an emitted token. `None` when `METRALE_GLM_MTP_HEAD_FP8=0`,
92 /// when the `gemv_fp8w` kernel is absent, or when quantisation fails.
93 /// Measured 2026-08-29 (nsys): 2.66 ms -> ~1.33 ms per draft sweep.
94 head_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
95 gemv_fp8w_k: KernelHandle,
96}
97
98/// 2026-09-25: Rows the GLM drafter can ever be asked for: the served context, clamped to
99/// the rows its DSA indexer cache reserves. `max_dsa_context` is `max_context` (the serve's
100/// `--max-seq-len`) rounded down to whole `index_kpool` pools, so the cap follows the config.
101fn drafter_context_rows(max_seq_len: usize, cfg: &crate::glm5next_dsa::Glm5NextDsaConfig) -> usize {
102 max_seq_len.min(crate::glm5next_dsa::state::max_dsa_context(cfg))
103}
104
105impl Glm5NextMtpHead {
106 fn norm(
107 &self,
108 gpu: &dyn GpuBackend,
109 x: DevicePtr,
110 w: DevicePtr,
111 out: DevicePtr,
112 n: usize,
113 stream: u64,
114 ) -> Result<()> {
115 KernelLaunch::new(gpu, self.rms_norm_k)
116 .grid([1, 1, 1])
117 .block([(n.min(1024)) as u32, 1, 1])
118 .arg_ptr(x)
119 .arg_ptr(w)
120 .arg_ptr(out)
121 .arg_u32(n as u32)
122 .arg_f32(self.module.layer.rms_eps)
123 .launch(stream)
124 }
125
126 /// 2026-09-25: One draft token. Advances `seq_len` by one.
127 fn forward_one(
128 &self,
129 token: u32,
130 hidden_in: DevicePtr,
131 position: usize,
132 st: &mut Glm5NextMtpProposerState,
133 ctx: &ForwardContext,
134 stream: u64,
135 ) -> Result<u32> {
136 let gpu = ctx.gpu;
137 let h = self.hidden;
138 if position >= self.max_seq_len {
139 bail!(
140 "GLM MTP drafter: position {position} is past the {} it was sized for",
141 self.max_seq_len
142 );
143 }
144 // 2026-09-25: The embedding row is read by pointer: `embed_tokens` is `[vocab, hidden]`
145 // BF16 and a row is contiguous, so there is nothing to gather.
146 let embed_row = self.embed_tokens.weight.offset(token as usize * h * 2);
147 self.norm(gpu, embed_row, self.module.enorm, st.concat, h, stream)?;
148 self.norm(
149 gpu,
150 hidden_in,
151 self.module.hnorm,
152 st.concat.offset(h * 2),
153 h,
154 stream,
155 )?;
156 ops::dense_gemv(
157 gpu,
158 self.gemv_k,
159 st.concat,
160 &self.module.eh_proj,
161 st.x,
162 h as u32,
163 (2 * h) as u32,
164 stream,
165 )?;
166
167 // 2026-09-25: The block writes its output back over `st.x` (plain residual, in place).
168 if skip_block() {
169 st.seq_len += 1;
170 } else {
171 let mut kv = self.kv_cache.lock();
172 let dsa_state: &mut dyn LayerState = &mut st.dsa;
173 self.module.layer.decode_one_for_drafter(
174 st.x,
175 dsa_state,
176 &mut kv,
177 st.seq_len,
178 &mut st.block_table,
179 ctx,
180 stream,
181 )?;
182 drop(kv);
183 st.seq_len += 1;
184 }
185
186 // 2026-09-25: Timing arm `METRALE_GLM_MTP_SKIP=head`: the head is skipped and the
187 // draft is token 0, so `propose` times the block alone. `=block` skips the block instead.
188 if skip_head() {
189 return Ok(0);
190 }
191 // 2026-09-25: `shared_head.norm`, then the target's own `lm_head`: the MTP layer carries
192 // no head of its own.
193 self.norm(gpu, st.x, self.module.final_norm, st.x, h, stream)?;
194 // 2026-09-25: Sharded only with a communicator and a vocab split across ranks;
195 // otherwise the sweep covers the whole vocab.
196 let sharded = ctx.comm.is_some() && self.head_n != self.vocab;
197 let (w, n, v0) = if sharded {
198 (
199 DenseWeight {
200 weight: self.lm_head.weight.offset(self.head_v0 * h * 2),
201 },
202 self.head_n,
203 self.head_v0,
204 )
205 } else {
206 (self.lm_head, self.vocab, 0)
207 };
208 // 2026-09-25: The FP8 copy covers only `[head_v0, head_v0 + head_n)`, so it serves the
209 // sharded sweep alone; an unsharded sweep reads the BF16 head.
210 match self.head_fp8.filter(|_| sharded && n == self.head_n) {
211 Some(q) => ops::dense_gemv_fp8w(
212 gpu,
213 self.gemv_fp8w_k,
214 st.x,
215 &q,
216 st.logits,
217 n as u32,
218 h as u32,
219 stream,
220 )?,
221 None => ops::dense_gemv(
222 gpu,
223 self.gemv_k,
224 st.x,
225 &w,
226 st.logits,
227 n as u32,
228 h as u32,
229 stream,
230 )?,
231 }
232 ops::argmax_bf16(gpu, self.argmax_k, st.logits, st.arg, n as u32, stream)?;
233 let mut out = [0u8; 4];
234 gpu.synchronize(stream)?;
235 gpu.copy_d2h(st.arg, &mut out)?;
236 let local = u32::from_le_bytes(out) as usize;
237 let Some(comm) = ctx.comm.filter(|_| sharded) else {
238 return Ok((v0 + local) as u32);
239 };
240 // 2026-09-25: Exchange (max, argmax) in 8 BF16 lanes: `[val_r0, val_r1, then 3 base-256
241 // digits of each rank's global index]`. Each rank writes only its own lanes and leaves
242 // the others zero, so a SUM all-reduce delivers both ranks' values untouched (`x + 0.0`
243 // is exact).
244 //
245 // The all-reduce sums its bytes as BF16 elements (`NcclDataType::Bfloat16`), so an f32
246 // packed here would be summed as two BF16 halves. A token id needs 18 bits and BF16
247 // holds integers exactly only through 256, hence the digits; the logit lane is already
248 // BF16, so it round-trips bit for bit.
249 let mut lb = [0u8; 2];
250 gpu.copy_d2h(st.logits.offset(local * 2), &mut lb)?;
251 let g = v0 + local;
252 let bf = |x: f32| ((x.to_bits() >> 16) as u16).to_le_bytes();
253 let mut pack = [0u8; 16];
254 pack[self.head_rank * 2..][..2].copy_from_slice(&lb);
255 for d in 0..3 {
256 let digit = ((g >> (8 * d)) & 0xFF) as f32;
257 pack[4 + (self.head_rank * 3 + d) * 2..][..2].copy_from_slice(&bf(digit));
258 }
259 gpu.copy_h2d(&pack, st.head_xchg)?;
260 comm.all_reduce_async(st.head_xchg.0, 16, stream)?;
261 gpu.synchronize(stream)?;
262 gpu.copy_d2h(st.head_xchg, &mut pack)?;
263 let lane = |i: usize| {
264 f32::from_bits(
265 (u16::from_le_bytes(pack[i * 2..][..2].try_into().unwrap()) as u32) << 16,
266 )
267 };
268 // 2026-09-25: `>=` makes the lower rank win a tie, identically on both ranks, so the
269 // two ranks' drafter KV streams cannot diverge on a tie.
270 let win = if lane(0) >= lane(1) { 0 } else { 1 };
271 let idx = (0..3).fold(0usize, |a, d| {
272 a + ((lane(2 + win * 3 + d) as usize) << (8 * d))
273 });
274 Ok(idx as u32)
275 }
276
277 /// 2026-09-25: Append `tokens.len() - 1` drafter context rows: row `r` is pair key
278 /// `row_base + r` = `(embed(tokens[r + 1]), hiddens row r)`. Used for both the whole-prompt
279 /// prefill and the catch-up feed; the only difference between them is `row_base`.
280 ///
281 /// The row space is dense: a row's KV slot, indexer row and RoPE position are all
282 /// `seq_len` (see `Glm5NextDsaLayer::write_kv_row`), so every pair key from 0 up is
283 /// written and slot == key.
284 ///
285 /// Unless `METRALE_GLM_MTP_PREFILL_FULL=1`, a context row runs only `input_norm` and the
286 /// DSA cache writes (`drafter_write_kv_row`): no MoE, no attention, no `lm_head`, because
287 /// its block output is discarded.
288 #[allow(clippy::too_many_arguments)]
289 fn rows_impl(
290 &self,
291 tokens: &[u32],
292 hiddens: DevicePtr,
293 row_base: usize,
294 state: &mut dyn ProposerState,
295 ctx: &ForwardContext,
296 stream: u64,
297 ) -> Result<usize> {
298 let st = match state
299 .as_any_mut()
300 .downcast_mut::<Glm5NextMtpProposerState>()
301 {
302 Some(s) => s,
303 None => return Ok(0),
304 };
305 // 2026-09-25: Rows must append exactly at the drafter's current length, or the dense
306 // row space grows a hole and every later RoPE position is wrong.
307 if st.seq_len != row_base || tokens.len() < 2 {
308 return Ok(0);
309 }
310 let h = self.hidden;
311 let rows = tokens.len() - 1;
312 if row_base + rows > self.max_seq_len {
313 return Ok(0);
314 }
315 let gpu = ctx.gpu;
316 let dbg = metrale_model_layers::speculative::mtp_refeed_debug();
317 let prefill_full = std::env::var("METRALE_GLM_MTP_PREFILL_FULL")
318 .ok()
319 .as_deref()
320 == Some("1");
321 let mut kv = self.kv_cache.lock();
322 let Glm5NextMtpProposerState {
323 dsa,
324 seq_len,
325 block_table,
326 concat,
327 x,
328 ..
329 } = st;
330 for r in 0..rows {
331 let embed_row = self
332 .embed_tokens
333 .weight
334 .offset(tokens[r + 1] as usize * h * 2);
335 self.norm(gpu, embed_row, self.module.enorm, *concat, h, stream)?;
336 self.norm(
337 gpu,
338 hiddens.offset(r * h * 2),
339 self.module.hnorm,
340 concat.offset(h * 2),
341 h,
342 stream,
343 )?;
344 ops::dense_gemv(
345 gpu,
346 self.gemv_k,
347 *concat,
348 &self.module.eh_proj,
349 *x,
350 h as u32,
351 (2 * h) as u32,
352 stream,
353 )?;
354 let dsa_state: &mut dyn LayerState = dsa;
355 // 2026-09-25: Diagnostic arm `METRALE_GLM_MTP_PREFILL_FULL=1`: build the row through
356 // the full block path a propose uses (`decode_one_for_drafter`) instead of the
357 // cache-only write.
358 if prefill_full {
359 self.module.layer.decode_one_for_drafter(
360 *x,
361 dsa_state,
362 &mut kv,
363 *seq_len,
364 block_table,
365 ctx,
366 stream,
367 )?;
368 } else {
369 self.module.layer.drafter_write_kv_row(
370 *x,
371 dsa_state,
372 &mut kv,
373 *seq_len,
374 block_table,
375 ctx,
376 stream,
377 )?;
378 }
379 if dbg {
380 let fp = metrale_model_layers::speculative::hidden_fingerprint(
381 gpu,
382 hiddens.offset(r * h * 2),
383 h,
384 );
385 tracing::info!(
386 "GLM_MTP_DBG ctx row slot={} key={} tok={} fp_hidden={fp:016x}",
387 *seq_len,
388 row_base + r,
389 tokens[r + 1],
390 );
391 }
392 *seq_len += 1;
393 }
394 Ok(rows)
395 }
396}
397
398/// 2026-09-25: `METRALE_GLM_MTP_SKIP=head`: stop the drafter after the block, before
399/// `shared_head.norm`, the `lm_head` gemv, the argmax and the D2H. A timing arm; read once
400/// per process.
401fn skip_head() -> bool {
402 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
403 *ON.get_or_init(|| std::env::var("METRALE_GLM_MTP_SKIP").ok().as_deref() == Some("head"))
404}
405
406/// 2026-09-25: `METRALE_GLM_MTP_SKIP=block`: skip the MTP block itself and run only the
407/// head. A timing arm; read once per process.
408fn skip_block() -> bool {
409 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
410 *ON.get_or_init(|| std::env::var("METRALE_GLM_MTP_SKIP").ok().as_deref() == Some("block"))
411}
412
413#[cfg(test)]
414#[path = "glm5next_mtp_head_tests.rs"]
415mod a59_sizing_tests;