metrale_model_arch/glm5next_mtp_head/
init.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Glm5NextMtpHead::new`: the drafter's private KV pool and its vocab-sharded draft head.
4//!
5//! Owner: model-arch (GLM-5.3 MTP drafter).
6//! Invariants:
7//! - A constructed head's `max_seq_len` is at most `max_dsa_context` of its DSA block.
8//! - `head_n` is `vocab_size / tp_world_size` when `tp_world_size > 1` divides
9//!   `vocab_size`, and `vocab_size` otherwise.
10
11use super::*;
12
13impl Glm5NextMtpHead {
14    pub fn new(
15        module: Glm5NextMtpModule,
16        embed_tokens: DenseWeight,
17        lm_head: DenseWeight,
18        config: &metrale_config::ModelConfig,
19        gpu: &dyn GpuBackend,
20        max_seq_len: usize,
21    ) -> Result<Self> {
22        let dsa = match &module.layer.mixer {
23            crate::glm5next_layer::Glm5NextMixer::Dsa(l) => l,
24            _ => bail!("GLM MTP block is not a DSA layer"),
25        };
26        // 2026-09-25: The drafter block is a DSA layer, so it can never reach a position past
27        // `max_dsa_context`: that is the indexer cache `Glm5NextDsaState::alloc` reserves, and
28        // `advance` refuses to grow past it. Everything sized from `max_seq_len` here uses the
29        // capped value: the private KV pool below, the block table claimed in `alloc_state`,
30        // the bounds checks in `forward_one` and `rows_impl`, and, through
31        // `prefill_hidden_rows`, the model's `mtp_prefill_hidden` capture.
32        let max_seq_len = drafter_context_rows(max_seq_len, &dsa.cfg);
33        // 2026-09-25: One KV head of `kv_lora_rank` per token, the absorbed-MLA latent that
34        // the block's `latent_write` and paged gather address.
35        let kv_config = KvCacheConfig {
36            block_size: 16,
37            num_kv_heads: 1,
38            head_dim: dsa.cfg.kv_lora_rank,
39            num_layers: 1,
40            dtype: KvCacheDtype::Fp8,
41            layer_dtypes: vec![],
42            layer_dims: vec![],
43            cache_blocks_per_seq: None,
44        };
45        let blocks = max_seq_len / kv_config.block_size + 2;
46        let kv_cache = PagedKvCache::new(kv_config, blocks, gpu)?;
47        // 2026-09-25: The draft `lm_head` sweep is most of the drafter's time. Measured
48        // 2026-08-29 with `METRALE_GLM_MTP_SKIP=head`: propose 8.84 ms -> 1.52 ms. It reads
49        // `[vocab, hidden]` BF16, 1.27 GB at GLM-5.3's 154,880 x 4,096.
50        //
51        // So the sweep is split by vocab across the ranks that run propose together: each
52        // rank sweeps its `head_n` rows over the full hidden dimension, and the ranks exchange
53        // (max, argmax) through one 16-byte all-reduce (`forward_one`).
54        //
55        // Vocab, not hidden: rows of `[vocab, hidden]` are contiguous, so a vocab shard is a
56        // base-pointer offset and a smaller `n`. A hidden shard would need a row stride, and
57        // the gemv takes none; it assumes rows packed at K.
58        let head_world = config.tp_world_size.max(1);
59        let head_rank = config.tp_rank;
60        let head_n = if head_world > 1 && config.vocab_size.is_multiple_of(head_world) {
61            config.vocab_size / head_world
62        } else {
63            config.vocab_size
64        };
65        // 2026-09-25: Quantise only the rows this rank sweeps: `head_n * hidden` bytes, not
66        // the whole vocab. A failure here is not fatal; the head falls back to the BF16 sweep.
67        let gemv_fp8w_k =
68            metrale_model_layers::layers::try_kernel(gpu, "gemv_fp8w", "dense_gemv_fp8w");
69        let head_fp8 = if std::env::var("METRALE_GLM_MTP_HEAD_FP8").as_deref() == Ok("0")
70            || gemv_fp8w_k.0 == 0
71        {
72            None
73        } else {
74            let shard = DenseWeight {
75                weight: lm_head
76                    .weight
77                    .offset(head_rank * head_n * config.hidden_size * 2),
78            };
79            match gpu
80                .kernel("gemv_fp8w", "quantize_bf16_to_fp8")
81                .and_then(|qk| {
82                    metrale_model_layers::weight_map::quantize_to_fp8(
83                        &shard,
84                        head_n,
85                        config.hidden_size,
86                        gpu,
87                        qk,
88                        gpu.default_stream(),
89                    )
90                }) {
91                Ok(q) => {
92                    tracing::info!(
93                        "GLM MTP: draft lm_head shard quantised to FP8 ({} rows x {}, {} MB)",
94                        head_n,
95                        config.hidden_size,
96                        head_n * config.hidden_size / (1024 * 1024),
97                    );
98                    Some(q)
99                }
100                Err(e) => {
101                    tracing::warn!("GLM MTP: FP8 draft head unavailable ({e:#}); staying BF16");
102                    None
103                }
104            }
105        };
106
107        Ok(Self {
108            module,
109            embed_tokens,
110            lm_head,
111            kv_cache: Mutex::new(kv_cache),
112            rms_norm_k: gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?,
113            gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?,
114            argmax_k: gpu.kernel("argmax", "argmax_bf16")?,
115            hidden: config.hidden_size,
116            vocab: config.vocab_size,
117            max_seq_len,
118            head_rank,
119            head_v0: head_rank * head_n,
120            head_n,
121            head_fp8,
122            gemv_fp8w_k,
123        })
124    }
125}
metrale_model_arch/glm5next_mtp_head/init.rs

metrale_model_arch/glm5next_mtp_head/
init.rs
