spark_model/weight_loader/
glm5_next_load.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `Glm5NextWeightLoader` — assembles the 45-layer GLM-5.3 text stack from a `WeightStore`.
4//!
5//! Everything it calls already existed and was gated: `bind_kda_weights`, `build_dsa_weights`,
6//! `glm5next_mlp::build`. This is the wiring, plus the one thing wiring must do that the pieces
7//! cannot — decide, per layer, WHICH pieces.
8//!
9//! # Where the classification comes from
10//!
11//! Not from tensor names, and not from modular arithmetic. `Glm5NextTextSkeleton::from_config`
12//! derives the mixer and MLP kind of all 45 layers from the checkpoint's own
13//! `linear_attn_config` index lists and `first_k_dense_replace`, cross-checked against the
14//! textual arrays, and refuses anything it was not taught. This loader iterates that.
15//!
16//! # TP, on every half
17//!
18//! DSA shards through `DsaTpPlan`, the MLP through `Glm5NextMlpConfig`, and KDA through
19//! [`KdaShardedSource`] — an adapter that slices the host bytes **before** the proven
20//! `bind_kda_weights` sees them, so TP=1 and TP=2 take the identical binder code path.
21//!
22//! 🪤 Both mixers end in a **row-parallel** `o_proj`, so the attention output is a partial sum
23//! at TP>1 and `Glm5NextLayer::mixer_all_reduce` reduces it before the mHC highway sees it.
24//! Half-applying the sharding — the state before this was wired — meant every rank computed a
25//! WHOLE KDA block and the all-reduce double-counted it: no crash, no shape error.
26
27use anyhow::{Context, Result, bail};
28use atlas_core::config::ModelConfig;
29use spark_runtime::gpu::{DevicePtr, GpuBackend};
30use spark_runtime::kv_cache::KvCacheDtype;
31use spark_runtime::weights::{WeightDtype, WeightStore, WeightTensor};
32
33use super::ModelWeightLoader;
34use crate::layer::TransformerLayer;
35use crate::layers::glm5next_dsa::build::build_dsa_weights;
36use crate::layers::glm5next_dsa::layer::{Glm5NextDsaLayer, Glm5NextDsaLayerKernels};
37use crate::layers::glm5next_dsa::{Glm5NextDsaConfig, Glm5NextDsaKernels};
38use crate::layers::glm5next_kda::binding::{
39    KdaDtype, KdaTensorSource, RawTensor, bind_kda_weights,
40};
41use crate::layers::glm5next_kda::tp::KdaTpPlan;
42use crate::layers::glm5next_kda::tp_bind::KdaShardedSource;
43use crate::layers::glm5next_kda::{Glm5NextKdaConfig, Glm5NextKdaKernels, Glm5NextKdaLayer};
44use crate::layers::glm5next_layer::{Glm5NextLayer, Glm5NextMhc, Glm5NextMixer, Glm5NextMlpSite};
45use crate::layers::glm5next_mlp::weights::{Glm5NextExpertWeights, Nvfp4Proj};
46use crate::layers::glm5next_mlp::{Glm5NextMlpConfig, Glm5NextMlpKernels, build as mlp_build};
47use crate::layers::glm5next_skeleton::{Glm5NextTextSkeleton, Mixer, Mlp};
48use crate::layers::ops::{Glm5NextMhcKernels, Glm5NextMhcSiteWeights, MHC_MIX_MAX_TOKENS, mix_hc};
49use crate::weight_map::DenseWeight;
50
51pub struct Glm5NextWeightLoader;
52
53/// A `[layer]`-relative tensor name, fully qualified for this checkpoint.
54///
55/// 🪤 GLM-5.3 nests the text stack under `model.language_model.`, not `model.`. And it does NOT
56/// use `mtp.0.*` — the MTP block is `layers.45`.
57fn qualify(layer: usize, leaf: &str) -> String {
58    format!("model.language_model.layers.{layer}.{leaf}")
59}
60
61/// Is this store tensor one the binders have already re-uploaded a copy of?
62///
63/// `load_layers` pulls every non-expert layer tensor to the host
64/// (`LayerSource::collect`) and the binders upload fresh device buffers — a TP
65/// shard for KDA/DSA, a dtype-converted copy for mHC. The store's originals are
66/// dead from that moment, and on GB10 they are 15.7 GB of unified memory the KV
67/// cache never gets. Measured 2026-08-28: the first 2-node bring-up died with
68/// "No memory left for KV cache" at 112.8 GB resident against a 99.64 GB load.
69///
70/// 🪤 Two things must NOT match:
71/// * `mlp.experts.*` — bound **zero-copy** from these very pointers
72///   (`bind_expert`). Freeing them is a use-after-free with no diagnostic.
73/// * `layers.{num_layers}` — GLM's MTP/draft block sits one past the skeleton
74///   (`layers.45` at `num_hidden_layers = 45`) and is read by
75///   `load_mtp_weights_multi`, not by `load_layers`.
76fn is_reuploaded(name: &str, num_layers: usize) -> bool {
77    let Some(rest) = name.strip_prefix("model.language_model.layers.") else {
78        return false;
79    };
80    let Some((idx, rel)) = rest.split_once('.') else {
81        return false;
82    };
83    let Ok(idx) = idx.parse::<usize>() else {
84        return false;
85    };
86    idx < num_layers && !rel.starts_with("mlp.experts.")
87}
88
89/// Read a device tensor back as host bytes.
90fn host_bytes(gpu: &dyn GpuBackend, t: &WeightTensor) -> Result<Vec<u8>> {
91    let mut b = vec![0u8; t.byte_size()];
92    gpu.copy_d2h(t.ptr, &mut b)?;
93    Ok(b)
94}
95
96/// Read a device tensor back as host `f32`, whatever width it is stored at.
97///
98/// 🪤 The dtype is read off the tensor, never assumed. `hc_*_fn` is BF16 on disk while the kernel
99/// wants F32, and `weight_scale_2` is F32 — this is the #341/#347 dtype-mismatch class, and
100/// the shapes never say so.
101fn host_f32(gpu: &dyn GpuBackend, t: &WeightTensor, what: &str) -> Result<Vec<f32>> {
102    let b = host_bytes(gpu, t)?;
103    match t.dtype {
104        WeightDtype::BF16 => Ok(b
105            .chunks_exact(2)
106            .map(|c| half::bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32())
107            .collect()),
108        WeightDtype::FP32 => Ok(b
109            .chunks_exact(4)
110            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
111            .collect()),
112        other => bail!(
113            "{what}: dtype {other:?} cannot be read as f32 without a conversion this loader refuses to guess"
114        ),
115    }
116}
117
118fn upload_f32(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
119    let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
120    let p = gpu.alloc(b.len().max(1))?;
121    gpu.copy_h2d(&b, p)?;
122    Ok(p)
123}
124
125/// One layer's slice of the store, presented to the KDA binder as layer-relative names.
126pub(super) struct LayerSource {
127    names: Vec<String>,
128    tensors: std::collections::BTreeMap<String, (WeightDtype, Vec<usize>, Vec<u8>)>,
129}
130
131impl LayerSource {
132    pub(super) fn collect(gpu: &dyn GpuBackend, store: &WeightStore, layer: usize) -> Result<Self> {
133        let prefix = format!("model.language_model.layers.{layer}.");
134        let mut names = Vec::new();
135        let mut tensors = std::collections::BTreeMap::new();
136        let rels: Vec<String> = store
137            .names()
138            .filter_map(|n| n.strip_prefix(&prefix).map(|r| r.to_string()))
139            .collect();
140        for rel in rels {
141            let rel = rel.as_str();
142            // The routed experts are the bulk of a layer and are bound zero-copy from their
143            // device pointers; pulling them to the host here would move gigabytes for nothing.
144            if rel.starts_with("mlp.experts.") {
145                names.push(rel.to_string());
146                continue;
147            }
148            names.push(rel.to_string());
149            let t = store.get(&format!("{prefix}{rel}"))?;
150            tensors.insert(
151                rel.to_string(),
152                (t.dtype, t.shape.clone(), host_bytes(gpu, t)?),
153            );
154        }
155        Ok(Self { names, tensors })
156    }
157
158    pub(super) fn f32(&self, name: &str) -> Result<Vec<f32>> {
159        let (dtype, _, bytes) = self
160            .tensors
161            .get(name)
162            .with_context(|| format!("missing tensor {name}"))?;
163        match dtype {
164            WeightDtype::BF16 => Ok(bytes
165                .chunks_exact(2)
166                .map(|c| half::bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32())
167                .collect()),
168            WeightDtype::FP32 => Ok(bytes
169                .chunks_exact(4)
170                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
171                .collect()),
172            other => bail!("{name}: dtype {other:?} is not a plain float tensor"),
173        }
174    }
175}
176
177impl KdaTensorSource for LayerSource {
178    fn get(&self, name: &str) -> Option<RawTensor<'_>> {
179        let (dtype, shape, bytes) = self.tensors.get(name)?;
180        // 🪤 A KDA block is entirely BF16 except `A_log`/`dt_bias`, which are F32. Anything
181        // else here is not a KDA tensor, and the binder must see the absence rather than a
182        // coerced dtype — it refuses on a dtype mismatch precisely because a cast would
183        // change the numerics.
184        let dtype = match dtype {
185            WeightDtype::BF16 => KdaDtype::Bf16,
186            WeightDtype::FP32 => KdaDtype::F32,
187            _ => return None,
188        };
189        Some(RawTensor {
190            dtype,
191            shape: shape.clone(),
192            bytes,
193        })
194    }
195    fn names(&self) -> Vec<String> {
196        self.names.clone()
197    }
198}
199
200/// Bind the two mHC sites of one layer.
201///
202/// 🪤 `hc_*_fn` is **BF16 on disk and F32 at the kernel**; `base`/`scale` are already F32. All
203/// three are uploaded as F32 here. Passing the on-disk BF16 straight through is exactly the
204/// defect class that produced #341 and #347.
205fn bind_mhc_site(
206    gpu: &dyn GpuBackend,
207    src: &LayerSource,
208    site: &str,
209    hc_mult: usize,
210    hidden: usize,
211) -> Result<Glm5NextMhcSiteWeights> {
212    let f = src.f32(&format!("hc_{site}_fn"))?;
213    let want = mix_hc(hc_mult) * hc_mult * hidden;
214    if f.len() != want {
215        bail!(
216            "hc_{site}_fn has {} elements, expected mix_hc({hc_mult}) * {hc_mult} * {hidden} = {want}",
217            f.len()
218        );
219    }
220    let base = src.f32(&format!("hc_{site}_base"))?;
221    if base.len() != mix_hc(hc_mult) {
222        bail!(
223            "hc_{site}_base has {} entries, expected mix_hc({hc_mult}) = {}",
224            base.len(),
225            mix_hc(hc_mult)
226        );
227    }
228    let scale = src.f32(&format!("hc_{site}_scale"))?;
229    if scale.len() != 3 {
230        bail!(
231            "hc_{site}_scale has {} entries, expected 3 (pre, post, comb)",
232            scale.len()
233        );
234    }
235    Ok(Glm5NextMhcSiteWeights {
236        // 🔴 BF16, because that is what the checkpoint stores (`[24, 16384]`, BF16 in the
237        // safetensors header). Uploading it as F32 doubled `hc_mix`'s traffic — 1.57 MB
238        // instead of 0.79 MB per site, 90 sites per token — for values that were already
239        // exactly BF16. `glm5next_hc_mix_bf16` reads it at that width and is bit-identical.
240        hc_fn: upload_f32_as_bf16(gpu, &f)?,
241        hc_fn_bf16: true,
242        hc_scale: upload_f32(gpu, &scale)?,
243        hc_base: upload_f32(gpu, &base)?,
244        // `hc_mix` -> `hc_finish` handoff. Per site so the layer's two sites cannot alias.
245        mix: gpu.alloc(MHC_MIX_MAX_TOKENS * mix_hc(hc_mult) * 4)?,
246    })
247}
248
249/// One routed expert, bound straight off the checkpoint's device pointers.
250pub(super) fn bind_expert(
251    gpu: &dyn GpuBackend,
252    store: &WeightStore,
253    layer: usize,
254    id: usize,
255) -> Result<Glm5NextExpertWeights> {
256    let proj = |p: &str| -> Result<Nvfp4Proj> {
257        let base = format!("mlp.experts.{id}.{p}");
258        let packed = store.get(&qualify(layer, &format!("{base}.weight")))?;
259        let scale = store.get(&qualify(layer, &format!("{base}.weight_scale")))?;
260        let s2 = store.get(&qualify(layer, &format!("{base}.weight_scale_2")))?;
261        if packed.dtype != WeightDtype::UInt8 {
262            bail!(
263                "{base}.weight is {:?}, expected packed U8 NVFP4",
264                packed.dtype
265            );
266        }
267        let s2 = host_f32(gpu, s2, &format!("{base}.weight_scale_2"))?;
268        let [s2] = s2[..] else {
269            bail!("{base}.weight_scale_2 is not a scalar");
270        };
271        Ok(Nvfp4Proj {
272            packed: packed.ptr,
273            scale: scale.ptr,
274            scale_2: s2,
275        })
276    };
277    Ok(Glm5NextExpertWeights {
278        gate_proj: proj("gate_proj")?,
279        up_proj: proj("up_proj")?,
280        down_proj: proj("down_proj")?,
281    })
282}
283
284fn dense(store: &WeightStore, name: &str) -> Result<DenseWeight> {
285    Ok(DenseWeight {
286        weight: store.get(name)?.ptr,
287    })
288}
289
290impl ModelWeightLoader for Glm5NextWeightLoader {
291    /// Text-only port. `weight_loader/glm5_next.rs` classifies `model.visual.*`
292    /// as `TensorRole::Vision` and excludes it from `is_required()`; nothing in
293    /// this loader binds it. Saying so here keeps the tower off the GPU in the
294    /// first place — on the LibertAIDAI NVFP4 checkpoint that is 1.05 GiB per
295    /// rank, sitting between `--speculative --num-drafts 2` and a serve that
296    /// fits (measured 2026-08-29: K=3 at 32 K needs 13.58 GiB against 12.07 free).
297    fn binds_vision_encoder(&self) -> bool {
298        false
299    }
300
301    /// All three halves shard: DSA by head, KDA by head/channel, the MLP by width (TP) and by
302    /// expert set (EP).
303    fn supports_tp(&self) -> bool {
304        true
305    }
306
307    fn load_layers(
308        &self,
309        store: &WeightStore,
310        config: &ModelConfig,
311        gpu: &dyn GpuBackend,
312        _layer_kv_dtypes: &[KvCacheDtype],
313    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
314        let skeleton = Glm5NextTextSkeleton::from_config(config)?;
315        // 🪤 `l2_eps` and `chunk` are NOT config keys — `l2_eps` is FLA's `1/sqrt(sum + eps)`
316        // convention and `chunk` is a prefill tiling width whose results are identical over
317        // 2..32. Everything else comes off the checkpoint, including `gate_lower_bound`, which
318        // the parser now refuses to default.
319        let kda_cfg = Glm5NextKdaConfig {
320            hidden: config.hidden_size,
321            heads: config.linear_num_value_heads,
322            head_dim: config.linear_value_head_dim,
323            conv_kernel: config.linear_conv_kernel_dim,
324            gate_lower_bound: config.linear_gate_lower_bound,
325            rms_norm_eps: config.rms_norm_eps as f32,
326            l2_eps: 1e-6,
327            chunk: 32,
328        };
329        kda_cfg.validate()?;
330        // 🪤 `gate_rank` is not a config key — it is `f_a_proj`'s row count, read off the
331        // checkpoint (128 on GLM-5.3). Reading it from layer 0 rather than assuming it means a
332        // checkpoint revision that changes the gate bottleneck fails loudly at load.
333        let gate_rank = {
334            let n = qualify(0, "self_attn.f_a_proj.weight");
335            let t = store.get(&n).with_context(|| {
336                format!("glm5_next: {n} is needed to size the KDA gate bottleneck")
337            })?;
338            *t.shape.first().context("f_a_proj has no rows")?
339        };
340        let kda_plan = KdaTpPlan::from_config(config, gate_rank)?;
341        let dsa_cfg = Glm5NextDsaConfig::from_config(config)?;
342        let mlp_cfg = Glm5NextMlpConfig::from_config(config)?;
343
344        let kda_kernels = Glm5NextKdaKernels::resolve(gpu)?;
345        let dsa_kernels = Glm5NextDsaKernels::resolve(gpu)?;
346        let dsa_layer_kernels = Glm5NextDsaLayerKernels::resolve(gpu)?;
347        let mlp_kernels = Glm5NextMlpKernels::resolve(gpu)?;
348        let mhc_kernels_probe = Glm5NextMhcKernels::resolve(gpu)?;
349        let rms_norm_k = gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?;
350        // Only the MTP layer's plain residual path uses this; a text layer's residual lives in
351        // the mHC highway. `try_kernel` so a target without it still serves the text stack.
352        let add_k = crate::layers::try_kernel(gpu, "bf16_add", "bf16_add_inplace");
353
354        // 🪤 All 34 KDA blocks have identical geometry, so ONE workspace serves them all.
355        //
356        // Sized for the widest speculative verify rather than one token. Prefill still runs
357        // token-by-token through `Glm5NextLayer::prefill` (the mHC highway forces per-token
358        // anyway), but `Glm5NextLayer::forward_k` sweeps the KDA weights ONCE for all K rows of
359        // a verify and needs `[K, ...]` scratch to do it. The cap is the batched GEMV kernel's
360        // own `MAX_M`: past it `ops::dense_mm_bf16` falls back to the tile GEMM, which is not
361        // bit-identical to the serial decode a verify must reproduce.
362        //
363        // Cost is a few MB for the whole model: the FP32 buffers are already sized to
364        // `t_pad = ceil(t / chunk) * chunk = 32` at t = 1, so only the BF16 `[t, *]` scratch
365        // grows.
366        // 🔴 The workspaces must also hold a batched PREFILL sub-chunk, which is wider than any
367        // verify (ANOMALIES A65 — `Glm5NextLayer::prefill` hands `forward_k` `PREFILL_ROWS`
368        // rows). Sizing to the verify width alone made `forward_k` bail the moment prefill
369        // used it. Cost is per-layer scratch that scales with rows, not with context.
370        // 🪤 `prefill_rows()` too, not just the constant: `ATLAS_GLM_PREFILL_ROWS` can widen the
371        // sub-chunk at launch, and a workspace built for the default would make `forward_k` bail
372        // the first time the A/B lever was actually used.
373        let verify_k = (crate::layers::ops::DENSE_GEMV_BATCHM_MAX_M as usize)
374            .max(crate::layers::glm5next_layer::PREFILL_ROWS)
375            .max(crate::layers::glm5next_layer::prefill_rows());
376        let kda_ws = std::sync::Arc::new(crate::layers::glm5next_kda::Glm5NextKdaWorkspace::new(
377            gpu, &kda_cfg, verify_k,
378        )?);
379
380        let dsa_plan = crate::layers::glm5next_dsa::tp::DsaTpPlan::new(
381            config.tp_rank,
382            config.tp_world_size.max(1),
383            &dsa_cfg,
384        )?;
385        let last = skeleton.layers.len() - 1;
386        let mut out: Vec<Box<dyn TransformerLayer>> = Vec::with_capacity(skeleton.layers.len());
387
388        // 🪤 The KV pool is sized to `num_attention_layers()` (11 on GLM-5.3 — the
389        // sparse blocks), so a DSA layer must address it by its ordinal among
390        // KV-consuming layers, NOT by its index in the 45-layer model stack. The
391        // 34 KDA blocks carry recurrent state and take no pool slot.
392        let mut attn_layer_idx = 0usize;
393
394        for sl in &skeleton.layers {
395            let idx = sl.index;
396            let src = LayerSource::collect(gpu, store, idx)
397                .with_context(|| format!("glm5_next: collecting layer {idx}"))?;
398
399            let mixer = match sl.mixer {
400                Mixer::Kda => {
401                    // The adapter yields THIS RANK's slice with local shapes; the binder
402                    // validates against the (already local) config exactly as at TP=1.
403                    let sharded = KdaShardedSource::new(&src, &kda_plan)?;
404                    let (w, _report) = bind_kda_weights(gpu, &kda_cfg, idx, &sharded)?;
405                    Glm5NextMixer::Kda {
406                        layer: Box::new(Glm5NextKdaLayer::new(idx, kda_cfg, w, kda_kernels)?),
407                        ws: kda_ws.clone(),
408                        cfg: kda_cfg,
409                    }
410                }
411                Mixer::Dsa => {
412                    let load = |n: &str| src.f32(n);
413                    let w = build_dsa_weights(gpu, &dsa_cfg, &dsa_plan, &load)?;
414                    Glm5NextMixer::Dsa(Box::new(Glm5NextDsaLayer {
415                        // ON by default since A55 was closed (the `weights_proj` overrun
416                        // fix). Kill switch `ATLAS_GLM_DSA_ALLOC_PER_STEP=1` restores the
417                        // per-step `gpu.alloc` + `gpu.free`.
418                        persist_bt: std::env::var("ATLAS_GLM_DSA_ALLOC_PER_STEP").as_deref()
419                            != Ok("1"),
420                        cfg: dsa_cfg,
421                        weights: w,
422                        kernels: dsa_layer_kernels,
423                        select_kernels: dsa_kernels,
424                        decode_kernel:
425                            crate::layers::glm5next_dsa::attend::Glm5NextDsaDecodeKernel::resolve(
426                                gpu,
427                            )?,
428                        workspace: crate::layers::glm5next_dsa::layer::Glm5NextDsaWorkspace::new(
429                            gpu, &dsa_cfg, verify_k,
430                        )?,
431                        layer_idx: idx,
432                        attn_layer_idx: {
433                            let a = attn_layer_idx;
434                            attn_layer_idx += 1;
435                            a
436                        },
437                        rms_eps: config.rms_norm_eps as f32,
438                        kv_scale: 1.0,
439                    }))
440                }
441            };
442
443            let load = |n: &str| src.f32(n);
444            let mlp = match sl.mlp {
445                Mlp::Dense => Glm5NextMlpSite::Dense(mlp_build::build_dense_mlp(
446                    gpu,
447                    &mlp_cfg,
448                    config.tp_rank,
449                    config.intermediate_size,
450                    "mlp",
451                    &load,
452                )?),
453                Mlp::RoutedMoe => {
454                    let expert = |id: usize| bind_expert(gpu, store, idx, id);
455                    Glm5NextMlpSite::Moe(Box::new(mlp_build::build_moe(
456                        gpu,
457                        &mlp_cfg,
458                        config.tp_rank,
459                        config.shared_expert_intermediate_size,
460                        &load,
461                        &expert,
462                    )?))
463                }
464            };
465
466            let mhc = if sl.hyper_connection {
467                Some(Glm5NextMhc {
468                    kernels: mhc_kernels_probe,
469                    attn: bind_mhc_site(gpu, &src, "attn", config.hc_mult, config.hidden_size)?,
470                    ffn: bind_mhc_site(gpu, &src, "ffn", config.hc_mult, config.hidden_size)?,
471                    hc_mult: config.hc_mult,
472                    sinkhorn_iters: config.hc_sinkhorn_iters,
473                    hc_eps: config.hc_eps,
474                })
475            } else {
476                None
477            };
478
479            out.push(Box::new(Glm5NextLayer {
480                layer_idx: idx,
481                mixer,
482                mlp,
483                mlp_cfg,
484                mlp_kernels,
485                mlp_ws: crate::layers::glm5next_mlp::forward::Glm5NextMlpWorkspace::new(
486                    gpu, &mlp_cfg, verify_k,
487                )?,
488                mhc,
489                input_norm: upload_f32_as_bf16(gpu, &src.f32("input_layernorm.weight")?)?,
490                post_attn_norm: upload_f32_as_bf16(
491                    gpu,
492                    &src.f32("post_attention_layernorm.weight")?,
493                )?,
494                rms_norm_k,
495                add_k,
496                rms_eps: config.rms_norm_eps as f32,
497                hidden: config.hidden_size,
498                mixer_all_reduce: match sl.mixer {
499                    Mixer::Kda => kda_plan.needs_output_all_reduce(),
500                    Mixer::Dsa => dsa_plan.needs_output_all_reduce(),
501                },
502                is_first: idx == 0,
503                is_last: idx == last,
504            }));
505        }
506        Ok(out)
507    }
508
509    fn load_embedding(
510        &self,
511        store: &WeightStore,
512        _config: &ModelConfig,
513        _gpu: &dyn GpuBackend,
514    ) -> Result<DenseWeight> {
515        dense(store, "model.language_model.embed_tokens.weight")
516    }
517
518    fn load_final_norm(
519        &self,
520        store: &WeightStore,
521        _config: &ModelConfig,
522        _gpu: &dyn GpuBackend,
523    ) -> Result<DenseWeight> {
524        dense(store, "model.language_model.norm.weight")
525    }
526
527    fn load_lm_head(
528        &self,
529        store: &WeightStore,
530        _config: &ModelConfig,
531        _gpu: &dyn GpuBackend,
532    ) -> Result<DenseWeight> {
533        dense(store, "lm_head.weight")
534    }
535
536    /// MTP is deliberately out of scope for this slice. `None` = "no speculative head", which
537    /// the scheduler already handles; it is not a silent skip of something wired.
538    fn load_mtp_weights(
539        &self,
540        _store: &WeightStore,
541        _config: &ModelConfig,
542        _gpu: &dyn GpuBackend,
543    ) -> Result<Option<crate::weight_loader::MtpWeights>> {
544        Ok(None)
545    }
546
547    /// Drop the store's copy of everything `load_layers` re-uploaded.
548    fn prune_after_load(
549        &self,
550        store: &mut WeightStore,
551        config: &ModelConfig,
552        gpu: &dyn GpuBackend,
553    ) -> Result<()> {
554        let n = config.num_hidden_layers;
555        let (count, bytes) = store.free_matching(gpu, |name| is_reuploaded(name, n))?;
556        tracing::info!(
557            "glm5_next: released {count} store tensors ({:.2} GB) already re-uploaded by the \
558             binders; routed experts and the MTP block kept",
559            bytes as f64 / 1e9,
560        );
561        Ok(())
562    }
563}
564
565pub(super) fn upload_f32_as_bf16(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
566    let b: Vec<u8> = v
567        .iter()
568        .flat_map(|x| half::bf16::from_f32(*x).to_le_bytes())
569        .collect();
570    let p = gpu.alloc(b.len().max(1))?;
571    gpu.copy_h2d(&b, p)?;
572    Ok(p)
573}
574
575#[cfg(test)]
576mod prune_tests {
577    use super::is_reuploaded;
578
579    #[test]
580    fn prunes_only_the_reuploaded_layer_tensors() {
581        let n = 45;
582        // Re-uploaded by the binders -> free.
583        assert!(is_reuploaded(
584            "model.language_model.layers.0.self_attn.q_proj.weight",
585            n
586        ));
587        assert!(is_reuploaded(
588            "model.language_model.layers.44.mlp.gate.weight",
589            n
590        ));
591        assert!(is_reuploaded(
592            "model.language_model.layers.3.hc_attn_fn.weight",
593            n
594        ));
595        // Bound zero-copy from the store -> use-after-free if freed.
596        assert!(!is_reuploaded(
597            "model.language_model.layers.7.mlp.experts.12.down_proj.weight",
598            n
599        ));
600        // MTP block, one past the skeleton -> read by load_mtp_weights_multi.
601        assert!(!is_reuploaded(
602            "model.language_model.layers.45.self_attn.q_proj.weight",
603            n
604        ));
605        // Not a layer tensor at all.
606        assert!(!is_reuploaded(
607            "model.language_model.embed_tokens.weight",
608            n
609        ));
610        assert!(!is_reuploaded("lm_head.weight", n));
611        // Malformed / non-numeric index is never a match.
612        assert!(!is_reuploaded("model.language_model.layers.x.foo", n));
613    }
614}
615
616/// [`LayerSource::collect`], named for the MTP loader's call site.
617pub(super) fn layer_source(
618    gpu: &dyn GpuBackend,
619    store: &WeightStore,
620    layer: usize,
621) -> Result<LayerSource> {
622    LayerSource::collect(gpu, store, layer)
623}
624
625/// [`bind_expert`], named for the MTP loader's call site.
626pub(super) fn bind_expert_at(
627    gpu: &dyn GpuBackend,
628    store: &WeightStore,
629    layer: usize,
630    id: usize,
631) -> Result<Glm5NextExpertWeights> {
632    bind_expert(gpu, store, layer, id)
633}
634
635/// [`upload_f32_as_bf16`], named for the MTP loader's call site.
636pub(super) fn upload_bf16(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
637    upload_f32_as_bf16(gpu, v)
638}
639
640#[cfg(test)]
641mod vision_capability_tests {
642    use super::Glm5NextWeightLoader;
643    use crate::weight_loader::ModelWeightLoader;
644
645    #[test]
646    fn glm5_next_declares_itself_text_only() {
647        // Mutation gate: flipping this to `true` re-loads 1.05 GiB/rank of
648        // vision tower that nothing binds, and K=3 stops fitting at 32 K.
649        assert!(
650            !Glm5NextWeightLoader.binds_vision_encoder(),
651            "GLM-5.3's port binds no vision encoder; saying otherwise makes the \
652             weight loader read the tower into unified memory for nothing"
653        );
654    }
655
656    #[test]
657    fn a_multimodal_loader_still_declares_true_by_default() {
658        // The trait default must stay "load everything" — a loader that never
659        // overrides this must never lose weights.
660        assert!(crate::weight_loader::qwen35::Qwen35WeightLoader.binds_vision_encoder());
661    }
662}
spark_model/weight_loader/glm5_next_load.rs

spark_model/weight_loader/
glm5_next_load.rs
