metrale_model_arch/weight_loader/
glm5_next_mtp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Load GLM-5.3's MTP block, `model.language_model.layers.{num_hidden_layers}`,
4//! as the body of the draft proposer.
5//!
6//! Owner: model-arch weight loader (GLM-5.3).
7//! Invariants:
8//! - The block is one DSA mixer and one routed MoE with `mhc: None`, the plain-residual path
9//!   of [`crate::glm5next_layer::Glm5NextLayer`].
10//! - Its DSA layer uses `attn_layer_idx = 0` of the drafter's own one-layer KV pool
11//!   (`glm5next_mtp_head`), never the target's.
12//! - `Glm5NextWeightLoader::prune_after_load` keeps this layer's tensors: `is_reuploaded`
13//!   matches only layer indices below `num_hidden_layers`.
14
15use anyhow::{Context, Result};
16use metrale_config::ModelConfig;
17use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
18use metrale_model_weights::weights::WeightStore;
19
20use crate::glm5next_dsa::build::build_dsa_weights;
21use crate::glm5next_dsa::layer::{Glm5NextDsaLayer, Glm5NextDsaLayerKernels, Glm5NextDsaWorkspace};
22use crate::glm5next_dsa::{Glm5NextDsaConfig, Glm5NextDsaKernels};
23use crate::glm5next_layer::{Glm5NextLayer, Glm5NextMixer, Glm5NextMlpSite};
24use crate::glm5next_mlp::{Glm5NextMlpConfig, Glm5NextMlpKernels, build as mlp_build};
25use metrale_model_layers::weight_map::DenseWeight;
26
27/// 2026-09-25: The MTP block plus the four tensors around it.
28pub struct Glm5NextMtpModule {
29    /// 2026-09-25: The block as a `mhc: None` layer: DSA mixer, routed MoE, plain residual.
30    pub layer: Glm5NextLayer,
31    /// 2026-09-25: `[hidden, 2 * hidden]`: projects `concat(enorm(embed), hnorm(hidden))`
32    /// back to `hidden` (`glm5next_mtp_head.rs`).
33    pub eh_proj: DenseWeight,
34    /// 2026-09-25: RMSNorm weights for the embedding half and the target-hidden half of the
35    /// concat.
36    pub enorm: DevicePtr,
37    pub hnorm: DevicePtr,
38    /// 2026-09-25: `shared_head.norm`, the final norm before the target's `lm_head`.
39    pub final_norm: DevicePtr,
40}
41
42/// 2026-09-25: Build the MTP block, or `None` when the checkpoint has no
43/// `layers.{num_hidden_layers}`.
44///
45/// Runs on every rank: `build_moe` binds only this rank's `local_expert_range()`, so each
46/// rank holds its own share of the routed experts, as in the text stack.
47pub fn load_glm5next_mtp_module(
48    store: &WeightStore,
49    config: &ModelConfig,
50    gpu: &dyn GpuBackend,
51) -> Result<Option<Glm5NextMtpModule>> {
52    let idx = config.num_hidden_layers;
53    let prefix = format!("model.language_model.layers.{idx}.");
54    if !store.names().any(|n| n.starts_with(&prefix)) {
55        return Ok(None);
56    }
57    let src = super::glm5_next_load::layer_source(gpu, store, idx)
58        .with_context(|| format!("glm5_next MTP: collecting layer {idx}"))?;
59    let load = |n: &str| src.f32(n);
60
61    let dsa_cfg = Glm5NextDsaConfig::from_config(config)?;
62    let mlp_cfg = Glm5NextMlpConfig::from_config(config)?;
63    let dsa_plan = crate::glm5next_dsa::tp::DsaTpPlan::new(
64        config.tp_rank,
65        config.tp_world_size.max(1),
66        &dsa_cfg,
67    )?;
68    let dsa_layer_kernels = Glm5NextDsaLayerKernels::resolve(gpu)?;
69    let dsa_kernels = Glm5NextDsaKernels::resolve(gpu)?;
70    let mlp_kernels = Glm5NextMlpKernels::resolve(gpu)?;
71
72    let mixer = Glm5NextMixer::Dsa(Box::new(Glm5NextDsaLayer {
73        persist_bt: std::env::var("METRALE_GLM_DSA_ALLOC_PER_STEP").as_deref() != Ok("1"),
74        cfg: dsa_cfg,
75        weights: build_dsa_weights(gpu, &dsa_cfg, &dsa_plan, &load)?,
76        kernels: dsa_layer_kernels,
77        select_kernels: dsa_kernels,
78        decode_kernel: crate::glm5next_dsa::attend::Glm5NextDsaDecodeKernel::resolve(gpu)?,
79        // 2026-09-25: Single-row workspace: the drafter runs this layer one row at a time.
80        workspace: Glm5NextDsaWorkspace::new(gpu, &dsa_cfg, 1)?,
81        layer_idx: idx,
82        // 2026-09-25: Sole consumer of its own one-layer KV pool.
83        attn_layer_idx: 0,
84        rms_eps: config.rms_norm_eps as f32,
85        kv_scale: 1.0,
86    }));
87
88    let expert = |id: usize| super::glm5_next_load::bind_expert_at(gpu, store, idx, id);
89    let mlp = Glm5NextMlpSite::Moe(Box::new(mlp_build::build_moe(
90        gpu,
91        &mlp_cfg,
92        config.tp_rank,
93        config.shared_expert_intermediate_size,
94        &load,
95        &expert,
96    )?));
97
98    let up =
99        |n: &str| -> Result<DevicePtr> { super::glm5_next_load::upload_bf16(gpu, &src.f32(n)?) };
100    Ok(Some(Glm5NextMtpModule {
101        layer: Glm5NextLayer {
102            layer_idx: idx,
103            mixer,
104            mlp,
105            mlp_cfg,
106            mlp_kernels,
107            // 2026-09-25: Its own one-row workspace, not the text stack's shared one, which
108            // is sized for a prefill sub-chunk.
109            mlp_ws: std::sync::Arc::new(crate::glm5next_mlp::forward::Glm5NextMlpWorkspace::new(
110                gpu, &mlp_cfg, 1,
111            )?),
112            // 2026-09-25: No hyper-connection: nothing here binds `hc_*` tensors, and
113            // `mhc: None` selects the plain residual path.
114            mhc: None,
115            input_norm: up("input_layernorm.weight")?,
116            post_attn_norm: up("post_attention_layernorm.weight")?,
117            rms_norm_k: gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?,
118            add_k: metrale_model_layers::layers::try_kernel(gpu, "bf16_add", "bf16_add_inplace"),
119            rms_eps: config.rms_norm_eps as f32,
120            hidden: config.hidden_size,
121            // 2026-09-25: Row-parallel `o_proj`, as in the DSA text layers.
122            mixer_all_reduce: dsa_plan.needs_output_all_reduce(),
123            // 2026-09-25: No mHC highway to expand or collapse.
124            is_first: false,
125            is_last: false,
126        },
127        eh_proj: DenseWeight {
128            weight: store.get(&format!("{prefix}eh_proj.weight"))?.ptr,
129        },
130        enorm: up("enorm.weight")?,
131        hnorm: up("hnorm.weight")?,
132        final_norm: up("shared_head.norm.weight")?,
133    }))
134}
metrale_model_arch/weight_loader/glm5_next_mtp.rs

metrale_model_arch/weight_loader/
glm5_next_mtp.rs
