metrale_model_arch/weight_loader/
qwen4_exp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen4ExpWeightLoader`, the weight loader for `qwen4_exp`
4//! (Qwen3.8-Flash-Next) checkpoints.
5//!
6//! Owner: model-arch weight loader.
7//! Invariants:
8//! - `load_layers` runs `audit_namespace` and `ensure_loadable` before it
9//!   builds any layer.
10//! - When `hc_mult > 0`, a load that succeeds has attached both mHC sites to
11//!   every layer.
12//!
13//! The GDN and full-attention layers are built by the qwen35 loader's arms
14//! (`linear_attn_arms`, `attention_arms`), keyed through
15//! `config.layer_prefix(i)`. On top of them this loader adds ones-filled
16//! placeholders for the per-layer norms, the multi-hyperconnection (mHC)
17//! sites (`hc`), the QSA indexer (`aux::attach_qsa`), the PLE layer (`ple`)
18//! and the MoE FFN (`ffn`).
19
20use anyhow::{Context, Result};
21use metrale_cache::kv_cache::KvCacheDtype;
22use metrale_config::{LayerType, ModelConfig};
23use metrale_gpu_runtime::gpu::GpuBackend;
24use metrale_model_weights::weights::WeightStore;
25
26use crate::weight_loader::ModelWeightLoader;
27use metrale_model_layers::layer::TransformerLayer;
28use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, dense};
29
30// 2026-09-25: The file is `aux_sites.rs` because `aux` is a reserved file
31// name on Windows, where a checkout of `aux.rs` fails.
32#[path = "qwen4_exp/aux_sites.rs"]
33mod aux;
34mod ffn;
35mod hc;
36mod ple;
37mod probe;
38
39pub use probe::audit_namespace;
40
41/// 2026-09-25: The PLE table's shard layout, `(file, byte offset)` per shard
42/// plus the rows per shard, read from a checkpoint's safetensors index and
43/// headers. `ple_tests.rs` uses it to build the segmented row cache without
44/// loading the model.
45#[cfg(test)]
46pub fn ple_shard_layout(snapshot: &str) -> Result<(Vec<(std::path::PathBuf, u64)>, u64)> {
47    use std::io::{Read, Seek, SeekFrom};
48    let idx: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(
49        std::path::Path::new(snapshot).join("model.safetensors.index.json"),
50    )?)?;
51    let map = idx["weight_map"].as_object().context("weight_map")?;
52    let mut names: Vec<(usize, &String)> = map
53        .keys()
54        .filter(|k| k.contains(".ngram_embedding.shard_"))
55        .map(|k| {
56            let n = k
57                .rsplit("shard_")
58                .next()
59                .and_then(|r| r.split('.').next())
60                .and_then(|r| r.parse().ok())
61                .unwrap_or(usize::MAX);
62            (n, k)
63        })
64        .collect();
65    names.sort();
66    anyhow::ensure!(!names.is_empty(), "no PLE shards in {snapshot}");
67
68    // 2026-09-25: Each file's header is read once. A shard's offset is taken
69    // against its own file's data start, because shards may be spread over
70    // several files.
71    let mut headers: std::collections::HashMap<String, (serde_json::Value, u64)> =
72        std::collections::HashMap::new();
73    let mut shards = Vec::with_capacity(names.len());
74    let mut rows_per = 0u64;
75    for (i, name) in &names {
76        let file = map[name.as_str()].as_str().context("shard file")?;
77        if !headers.contains_key(file) {
78            let path = std::path::Path::new(snapshot).join(file);
79            let mut fh = std::fs::File::open(&path)?;
80            let mut len = [0u8; 8];
81            fh.read_exact(&mut len)?;
82            let hlen = u64::from_le_bytes(len);
83            let mut hdr = vec![0u8; hlen as usize];
84            fh.seek(SeekFrom::Start(8))?;
85            fh.read_exact(&mut hdr)?;
86            headers.insert(file.to_owned(), (serde_json::from_slice(&hdr)?, 8 + hlen));
87        }
88        let (hdr, data_start) = &headers[file];
89        let e = &hdr[name.as_str()];
90        let off = e["data_offsets"][0].as_u64().context("data_offsets")?;
91        let rows = e["shape"][0].as_u64().context("shape")?;
92        if *i == 0 {
93            rows_per = rows;
94        }
95        anyhow::ensure!(
96            rows == rows_per,
97            "shard {i} has {rows} rows, not {rows_per}"
98        );
99        shards.push((std::path::Path::new(snapshot).join(file), data_start + off));
100    }
101    Ok((shards, rows_per))
102}
103
104pub struct Qwen4ExpWeightLoader;
105
106impl ModelWeightLoader for Qwen4ExpWeightLoader {
107    fn supports_tp(&self) -> bool {
108        false
109    }
110
111    fn load_layers(
112        &self,
113        store: &WeightStore,
114        config: &ModelConfig,
115        gpu: &dyn GpuBackend,
116        layer_kv_dtypes: &[KvCacheDtype],
117    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
118        let report = audit_namespace(store, config);
119        report.log();
120        report.ensure_loadable()?;
121
122        let h = config.hidden_size;
123        let variant = metrale_model_layers::weight_map::detect_nvfp4_variant(store, config);
124        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
125        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
126        let stream = gpu.default_stream();
127
128        tracing::info!(
129            "Qwen3.8-Flash-Next: {} layers ({} GDN + {} full attention), \
130             {} experts top-{}, hc {} streams x rank {}, indexer budget {}, \
131             PLE at {:?}; NVFP4 variant {:?}",
132            config.num_hidden_layers,
133            config
134                .layer_types
135                .iter()
136                .filter(|t| **t == LayerType::LinearAttention)
137                .count(),
138            config
139                .layer_types
140                .iter()
141                .filter(|t| **t == LayerType::FullAttention)
142                .count(),
143            config.num_experts,
144            config.num_experts_per_tok,
145            config.hc_mult,
146            config.hc_lowrank,
147            config.index_topk,
148            config.ple_layer_ids,
149            variant,
150        );
151
152        // 2026-09-25: The model-level mixer, which collapses the streams before
153        // `lm_head`. Every layer gets a copy; only the last model layer uses it.
154        let hc_head = if config.hc_mult > 0 {
155            Some(hc::load_head(store, config)?)
156        } else {
157            None
158        };
159
160        // 2026-09-25: The PLE scratch is sized once, for the largest prefill
161        // chunk, not the context: `METRALE_PLE_MAX_TOKENS`, default 2048. A
162        // larger chunk is refused by the PLE layer with a message naming the
163        // variable (`ple/layer.rs`).
164        let max_ple_tokens: usize = std::env::var("METRALE_PLE_MAX_TOKENS")
165            .ok()
166            .and_then(|v| v.parse().ok())
167            .unwrap_or(2048);
168        // 2026-09-25: With `METRALE_QWEN4EXP_NO_PLE=1`, `ple::load` is not
169        // called, so no row cache is opened.
170        let ple_off = std::env::var("METRALE_QWEN4EXP_NO_PLE").as_deref() == Ok("1");
171        // 2026-09-25: GDN layers are built with BF16 projections
172        // (`build_linear_attention_dense_bf16`) unless
173        // `METRALE_QWEN4EXP_BF16_GDN=0`, which requantizes them to NVFP4.
174        let bf16_gdn = std::env::var("METRALE_QWEN4EXP_BF16_GDN").as_deref() != Ok("0");
175        tracing::info!(
176            "GDN projections: {} on the {} linear-attention layers",
177            if bf16_gdn {
178                "BF16 as shipped (no runtime NVFP4 requantization)"
179            } else {
180                "requantized to NVFP4 (METRALE_QWEN4EXP_BF16_GDN=0)"
181            },
182            config
183                .layer_types
184                .iter()
185                .filter(|t| **t == LayerType::LinearAttention)
186                .count(),
187        );
188
189        let mut layers: Vec<Box<dyn TransformerLayer>> =
190            Vec::with_capacity(config.num_hidden_layers);
191        let mut attn_idx = 0usize;
192
193        // 2026-09-25: Free-memory deltas per part (MoE, attention/GDN arm,
194        // mHC+PLE), summed over layers and logged once.
195        let (mut moe_bytes, mut arm_bytes, mut hc_bytes) = (0u64, 0u64, 0u64);
196        let free_now = |g: &dyn GpuBackend| g.free_memory().unwrap_or(0) as u64;
197
198        for i in 0..config.num_hidden_layers {
199            let lp = config.layer_prefix(i);
200            let f0 = free_now(gpu);
201            let ffn = ffn::build_moe(store, &lp, config, gpu, variant)?;
202            let f1 = free_now(gpu);
203            moe_bytes += f0.saturating_sub(f1);
204
205            // 2026-09-25: Ones-filled per-layer norms for the shared arms'
206            // signatures. `ensure_loadable` refuses a checkpoint that has
207            // per-layer norms, and the low-rank mHC forward paths do not apply
208            // these (`HcVariant::applies_block_input_norm`,
209            // `qwen3_ssm/trait_decode_hc.rs`).
210            let input_norm = ones_norm(h, gpu)?;
211            let post_attn_norm = ones_norm(h, gpu)?;
212
213            let layer = match config.layer_types[i] {
214                LayerType::LinearAttention if bf16_gdn => {
215                    crate::weight_loader::qwen35::load_layers::linear_attn_arms::build_linear_attention_dense_bf16(
216                        i, store, &lp, gpu, variant, config, h,
217                        input_norm, post_attn_norm, ffn,
218                    )
219                    .with_context(|| format!("qwen4_exp: GDN layer {i} (BF16)"))?
220                }
221                LayerType::LinearAttention => {
222                    crate::weight_loader::qwen35::load_layers::linear_attn_arms::build_linear_attention_nvfp4(
223                        store, &lp, gpu, variant, config, h, absmax_k, quantize_k, stream,
224                        input_norm, post_attn_norm, ffn,
225                    )
226                    .with_context(|| format!("qwen4_exp: GDN layer {i}"))?
227                }
228                LayerType::FullAttention => {
229                    let kv_dtype = layer_kv_dtypes
230                        .get(attn_idx)
231                        .copied()
232                        .unwrap_or(KvCacheDtype::Bf16);
233                    let l = crate::weight_loader::qwen35::load_layers::attention_arms::build_full_attention_nvfp4(
234                        i, store, &lp, gpu, variant, config, h, absmax_k, quantize_k, stream,
235                        kv_dtype, attn_idx, input_norm, post_attn_norm, ffn,
236                    )
237                    .with_context(|| format!("qwen4_exp: full-attention layer {i}"))?;
238                    attn_idx += 1;
239                    l
240                }
241                other => anyhow::bail!(
242                    "qwen4_exp layer {i} has type {other:?}; this architecture is \
243                     only linear_attention / full_attention"
244                ),
245            };
246            let f2 = free_now(gpu);
247            arm_bytes += f1.saturating_sub(f2);
248
249            // 2026-09-25: Two mHC sites per layer, around attention and around
250            // the MoE.
251            let mut layer = layer;
252            if config.hc_mult > 0 {
253                let (attn, ffn) = hc::load_layer_sites(store, &lp, config)?;
254                aux::attach_hc(&mut layer, i, attn, ffn, hc_head.clone(), config)?;
255            }
256            aux::attach_qsa(&mut layer, i, &lp, store, config, gpu)?;
257            // 2026-09-25: `ple::load` returns `None` for a layer that
258            // `ple_layer_ids` does not list; `attach_ple` refuses a non-GDN host.
259            let ple_layer = if ple_off {
260                None
261            } else {
262                ple::load(store, config, i, max_ple_tokens, gpu)?
263            };
264            if let Some(p) = ple_layer {
265                aux::attach_ple(&mut layer, i, p)?;
266            }
267            layers.push(layer);
268            hc_bytes += f2.saturating_sub(free_now(gpu));
269        }
270        tracing::info!(
271            "qwen4_exp layer construction: MoE {:.2} GB ({:.1} MB/layer), \
272             attn/GDN arms {:.2} GB ({:.1} MB/layer), mHC+PLE {:.2} GB",
273            moe_bytes as f64 / 1e9,
274            moe_bytes as f64 / 1e6 / config.num_hidden_layers as f64,
275            arm_bytes as f64 / 1e9,
276            arm_bytes as f64 / 1e6 / config.num_hidden_layers as f64,
277            hc_bytes as f64 / 1e9,
278        );
279
280        if !config.ple_layer_ids.is_empty()
281            && std::env::var("METRALE_QWEN4EXP_NO_PLE").as_deref() == Ok("1")
282        {
283            tracing::warn!(
284                "METRALE_QWEN4EXP_NO_PLE=1: PLE n-gram injection at model layer {} \
285                 is DISABLED. Output is wrong by construction — this arm exists \
286                 to bisect the mHC spine, nothing else.",
287                config.ple_layer_ids[0].saturating_sub(1),
288            );
289        }
290        tracing::info!(
291            "Qwen3.8-Flash-Next loaded {} layers with the mHC highway live on \
292             all of them ({} GDN + {} full-attention).",
293            layers.len(),
294            layers.len()
295                - config
296                    .layer_types
297                    .iter()
298                    .filter(|t| **t == LayerType::FullAttention)
299                    .count(),
300            config
301                .layer_types
302                .iter()
303                .filter(|t| **t == LayerType::FullAttention)
304                .count(),
305        );
306        Ok(layers)
307    }
308
309    fn load_embedding(
310        &self,
311        store: &WeightStore,
312        config: &ModelConfig,
313        _gpu: &dyn GpuBackend,
314    ) -> Result<DenseWeight> {
315        let pfx = embed_prefix(config);
316        dense(store, &format!("{pfx}.embed_tokens.weight")).context("qwen4_exp: embedding")
317    }
318
319    /// 2026-09-25: A ones-filled `[hidden]` placeholder
320    /// (`aux::final_norm_placeholder`), which errors unless the mixer's
321    /// `hc_norm` exists. The model does not apply it: the qwen4_exp config
322    /// parser sets `final_norm_identity`, and the final-norm step then copies
323    /// the hidden state (`impl_a3_norm.rs`).
324    fn load_final_norm(
325        &self,
326        store: &WeightStore,
327        config: &ModelConfig,
328        gpu: &dyn GpuBackend,
329    ) -> Result<DenseWeight> {
330        aux::final_norm_placeholder(store, config, gpu)
331    }
332
333    fn load_lm_head(
334        &self,
335        store: &WeightStore,
336        config: &ModelConfig,
337        _gpu: &dyn GpuBackend,
338    ) -> Result<DenseWeight> {
339        if store.contains("lm_head.weight") {
340            return dense(store, "lm_head.weight");
341        }
342        anyhow::ensure!(
343            config.tie_word_embeddings,
344            "qwen4_exp: no lm_head.weight and tie_word_embeddings is false"
345        );
346        let pfx = embed_prefix(config);
347        dense(store, &format!("{pfx}.embed_tokens.weight")).context("qwen4_exp: tied lm_head")
348    }
349
350    fn load_vision_encoder(
351        &self,
352        store: &WeightStore,
353        config: &ModelConfig,
354        gpu: &dyn GpuBackend,
355    ) -> Result<Option<metrale_model_layers::layers::VisionTower>> {
356        crate::weight_loader::qwen35::Qwen35WeightLoader.load_vision_encoder(store, config, gpu)
357    }
358
359    fn load_mtp_weights(
360        &self,
361        _store: &WeightStore,
362        _config: &ModelConfig,
363        _gpu: &dyn GpuBackend,
364    ) -> Result<Option<MtpWeights>> {
365        Ok(None)
366    }
367}
368
369/// 2026-09-25: A ones-filled `[n]` BF16 norm scale. BF16 1.0 is `0x3F80`, two
370/// different bytes, so a byte `memset` cannot produce it.
371fn ones_norm(n: usize, gpu: &dyn GpuBackend) -> Result<DenseWeight> {
372    let host: Vec<u8> = std::iter::repeat_n([0x80u8, 0x3Fu8], n).flatten().collect();
373    let ptr = gpu.alloc(host.len())?;
374    gpu.copy_h2d(&host, ptr)?;
375    Ok(DenseWeight { weight: ptr })
376}
377
378/// 2026-09-25: `config.weight_prefix`, or `model` when it is empty.
379fn embed_prefix(config: &ModelConfig) -> String {
380    if config.weight_prefix.is_empty() {
381        "model".to_string()
382    } else {
383        config.weight_prefix.clone()
384    }
385}
386
387/// 2026-09-25: The prefix of the model-level hyper-connection mixer.
388fn mixer_prefix(config: &ModelConfig) -> String {
389    format!("{}.hyper_connection_mixer", embed_prefix(config))
390}
metrale_model_arch/weight_loader/qwen4_exp.rs

metrale_model_arch/weight_loader/
qwen4_exp.rs
