metrale_model_arch/weight_loader/
qwen35_dense.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen35DenseWeightLoader`, the loader for dense Qwen3.5-family checkpoints
4//! (`is_qwen35_dense`): full-attention and linear-attention (GDN) layers with a dense FFN,
5//! each built from the checkpoint's form (native FP8, NVFP4, BF16 or keep-packed Q2_0),
6//! and `prune_after_load`.
7//!
8//! Owner: model-arch weight loader (Qwen3.5 dense).
9//! Invariants:
10//! - `prune_after_load` frees a layer's store tensors only when the predicate
11//!   `load_layers` used for that layer (`gdn_fp8_arm_selected`,
12//!   `ffn_gateup_fused_selected`) holds.
13
14use anyhow::Result;
15use metrale_cache::kv_cache::KvCacheDtype;
16use metrale_config::{LayerType, ModelConfig};
17use metrale_gpu_runtime::gpu::GpuBackend;
18use metrale_model_weights::weights::{WeightDtype, WeightStore};
19
20use super::{ModelWeightLoader, WeightFormat};
21use metrale_model_layers::layer::TransformerLayer;
22use metrale_model_layers::weight_map::loaders_moe::load_mtp;
23use metrale_model_layers::weight_map::{
24    DenseWeight, Fp8Weight, MtpWeights, Nvfp4Variant, PackedQ2Weight, dense, detect_nvfp4_variant,
25};
26
27/// 2026-09-25: The Q2_0 group size when `{prefix}.weight` is a keep-packed ternary tensor
28/// (`WeightDtype::PackedQ2_0`, which the GGUF loader produces under
29/// `METRALE_GGUF_NATIVE_Q2=1`), else `None`.
30fn proj_q2_group(store: &WeightStore, prefix: &str) -> Option<u16> {
31    store
32        .get(&format!("{prefix}.weight"))
33        .ok()
34        .and_then(|w| w.q2_group())
35}
36
37/// 2026-09-25: Wraps the store's keep-packed `block_q2_0` buffer for `{prefix}.weight` as a
38/// [`PackedQ2Weight`] (`[n, k]` and group) without copying; the store keeps owning it.
39/// Errors unless the tensor is keep-packed Q2_0 and 2-D.
40fn packed_q2_from_store(store: &WeightStore, prefix: &str) -> Result<PackedQ2Weight> {
41    let w = store.get(&format!("{prefix}.weight"))?;
42    let group = w
43        .q2_group()
44        .ok_or_else(|| anyhow::anyhow!("{prefix}.weight is not keep-packed Q2_0"))?;
45    anyhow::ensure!(
46        w.shape.len() == 2,
47        "packed Q2_0 {prefix}.weight must be 2D, got {:?}",
48        w.shape
49    );
50    Ok(PackedQ2Weight {
51        weight: w.ptr,
52        n: w.shape[0] as u32,
53        k: w.shape[1] as u32,
54        group,
55    })
56}
57
58/// 2026-09-25: True when `{prefix}.weight` is FP8E4M3 with a block scale: a
59/// `.weight_scale_inv`, or a 2D `.weight_scale`. The same test as
60/// `qwen35::load_layers::proj_is_native_fp8`.
61fn proj_is_native_fp8(store: &WeightStore, prefix: &str) -> bool {
62    let is_fp8_weight = store
63        .get(&format!("{prefix}.weight"))
64        .map(|w| w.dtype == WeightDtype::FP8E4M3)
65        .unwrap_or(false);
66    let has_block_scale = store.contains(&format!("{prefix}.weight_scale_inv"))
67        || store
68            .get(&format!("{prefix}.weight_scale"))
69            .map(|s| s.shape.len() == 2)
70            .unwrap_or(false);
71    is_fp8_weight && has_block_scale
72}
73
74/// 2026-09-25: True when `{prefix}.weight` is a 2-D FP8E4M3 tensor whose
75/// `weight_scale_inv` or `weight_scale` is a single value or exactly the
76/// `[ceil(N/128), ceil(K/128)]` grid `w8a16` indexes as `block_scale[n/128, k/128]`
77/// (`load_fp8_block_scaled_as_fp8weight` repeats a single value over that grid).
78/// A per-row `[N, 1]` scale returns false: read as a grid, row `n` would take the scale
79/// in cell `n/128`. Such a projection is dequantized instead
80/// (`dequant_fp8_blockscaled_to_bf16` infers `[1, K]` blocks from the scale's shape).
81pub(super) fn proj_is_fp8_any_scale(store: &WeightStore, prefix: &str) -> bool {
82    let Ok(w) = store.get(&format!("{prefix}.weight")) else {
83        return false;
84    };
85    if w.dtype != WeightDtype::FP8E4M3 || w.shape.len() != 2 {
86        return false;
87    }
88    let (n, k) = (w.shape[0], w.shape[1]);
89
90    for key in [
91        format!("{prefix}.weight_scale_inv"),
92        format!("{prefix}.weight_scale"),
93    ] {
94        let Ok(s) = store.get(&key) else { continue };
95        if s.num_elements() == 1 {
96            return true;
97        }
98        if s.shape.len() == 2 && s.shape[0] == n.div_ceil(128) && s.shape[1] == k.div_ceil(128) {
99            return true;
100        }
101    }
102    false
103}
104
105/// 2026-09-25: Concatenates two block-scaled FP8 weights along rows,
106/// `[n_a, k] ++ [n_b, k] -> [n_a+n_b, k]`, into new weight and FP32 scale-grid buffers.
107/// The caller must ensure `n_a % 128 == 0`, so the two scale grids meet at a block
108/// boundary; nothing checks it.
109fn concat_fp8_block_scaled(
110    a: &Fp8Weight,
111    b: &Fp8Weight,
112    k: usize,
113    gpu: &dyn GpuBackend,
114) -> Result<Fp8Weight> {
115    let kb = k.div_ceil(128);
116    let a_w = a.n as usize * k;
117    let b_w = b.n as usize * k;
118    let weight = gpu.alloc(a_w + b_w)?;
119    gpu.copy_d2d(a.weight, weight, a_w)?;
120    gpu.copy_d2d(b.weight, weight.offset(a_w), b_w)?;
121    let a_s = (a.n as usize).div_ceil(128) * kb * 4;
122    let b_s = (b.n as usize).div_ceil(128) * kb * 4;
123    let row_scale = gpu.alloc(a_s + b_s)?;
124    gpu.copy_d2d(a.row_scale, row_scale, a_s)?;
125    gpu.copy_d2d(b.row_scale, row_scale.offset(a_s), b_s)?;
126    Ok(Fp8Weight {
127        weight,
128        row_scale,
129        n: a.n + b.n,
130        k: k as u32,
131        scale_format: metrale_model_layers::weight_map::WeightQuantFormat::Fp8BlockScaled,
132    })
133}
134
135/// 2026-09-25: `METRALE_DENSE_FP8=1` enables the native FP8 dense-FFN and attention arms
136/// (`ffn_fp8_arm_selected`, `attn_fp8`); off by default. Measured 2026-06-29 on a 9B FP8
137/// checkpoint: correct text at ~30 tok/s against ~40 tok/s on the NVFP4 path.
138fn dense_fp8_enabled() -> bool {
139    std::env::var("METRALE_DENSE_FP8").as_deref() == Ok("1")
140}
141
142/// 2026-09-25: Whether the native block-scaled FP8 GDN arm runs for this SSM layer.
143/// `load_layers` and `prune_after_load` both call it; the second frees the store tensors
144/// the first copied into the fused `[QKV|Z]` weight. The keep-packed Q2 arm is tested
145/// first in `load_layers`, so its condition is part of this one.
146fn gdn_fp8_arm_selected(store: &WeightStore, la: &str, tp_size: usize) -> bool {
147    let q2 = tp_size.max(1) == 1
148        && std::env::var_os("METRALE_NO_Q2_GDN").is_none()
149        && proj_q2_group(store, &format!("{la}.in_proj_qkv")).is_some()
150        && proj_q2_group(store, &format!("{la}.in_proj_z")).is_some();
151    !q2 && std::env::var_os("METRALE_NO_GDN_FP8").is_none()
152        && proj_is_fp8_any_scale(store, &format!("{la}.in_proj_qkv"))
153        && proj_is_fp8_any_scale(store, &format!("{la}.in_proj_z"))
154        && proj_is_fp8_any_scale(store, &format!("{la}.out_proj"))
155}
156
157/// 2026-09-25: The dense-FFN width: `intermediate_size`, or `moe_intermediate_size` when
158/// that is 0; the same order as `load_dense_ffn` (`weight_map/fp8_lut.rs`).
159fn ffn_inter(config: &ModelConfig) -> usize {
160    if config.intermediate_size > 0 {
161        config.intermediate_size
162    } else {
163        config.moe_intermediate_size
164    }
165}
166
167/// 2026-09-25: Whether the native FP8 dense-FFN arm runs for this layer:
168/// `METRALE_DENSE_FP8=1`, TP=1, the `Fp8Dequanted` variant and a block-scaled FP8
169/// `mlp.gate_proj`. `load_layers` calls it, and `prune_after_load` through
170/// `ffn_gateup_fused_selected`.
171fn ffn_fp8_arm_selected(
172    store: &WeightStore,
173    config: &ModelConfig,
174    variant: Nvfp4Variant,
175    lp: &str,
176) -> bool {
177    dense_fp8_enabled()
178        && config.tp_world_size.max(1) == 1
179        && matches!(variant, Nvfp4Variant::Fp8Dequanted)
180        && proj_is_native_fp8(store, &format!("{lp}.mlp.gate_proj"))
181}
182
183/// 2026-09-25: Whether this layer's gate and up are fused into one `[2*inter, hidden]`
184/// block-scaled FP8 weight: the FP8 FFN arm is selected, the target arms
185/// `ffn_gateup_fused` (`[defaults]`, or `METRALE_FFN_GATEUP_FUSED`), the model has no
186/// experts, both widths are multiples of 128 (so the concatenated scale grids meet on a
187/// block boundary), and both tensors are `[inter, hidden]`. `prune_after_load` frees
188/// the source tensors of exactly the layers this returns true for.
189fn ffn_gateup_fused_selected(
190    store: &WeightStore,
191    config: &ModelConfig,
192    variant: Nvfp4Variant,
193    lp: &str,
194) -> bool {
195    let inter = ffn_inter(config);
196    let hidden = config.hidden_size;
197    // 2026-09-25: Checked here rather than at the call site because `prune_after_load`
198    // frees on this answer: tensors whose shape is not `[inter, hidden]` from `config`
199    // are not fused, and so are not freed.
200    let on_disk = |name: &str| {
201        store
202            .get(&format!("{lp}.mlp.{name}.weight"))
203            .is_ok_and(|w| w.shape == [inter, hidden])
204    };
205    ffn_fp8_arm_selected(store, config, variant, lp)
206        && metrale_model_layers::layers::dense_ffn::gateup_fused::ffn_gateup_fused()
207        && config.num_experts == 0
208        && inter > 0
209        && inter.is_multiple_of(128)
210        && hidden.is_multiple_of(128)
211        && on_disk("gate_proj")
212        && on_disk("up_proj")
213}
214
215// 2026-09-25: `pub` because `weight_loader/mod.rs` re-exports them for the server's
216// preflight (`preflight/headroom.rs`), which reads `predicted_residency`; that module
217// prices its terms with `fp8_residency`'s shape helpers.
218pub mod fp8_residency;
219mod loaders_b;
220pub mod predicted_residency;
221mod rowwise_fp8;
222
223mod attn_arms;
224mod attn_layer;
225mod ffn_arm;
226mod gdn_dequant;
227mod gdn_layer;
228mod load_cx;
229mod prune;
230mod served_formats;
231mod w8a8_install;
232
233use fp8_residency::{DerivedResidency, RouteEnv};
234use load_cx::{Flow, LayerIn, LoadCx};
235
236pub struct Qwen35DenseWeightLoader;
237
238impl ModelWeightLoader for Qwen35DenseWeightLoader {
239    fn supports_tp(&self) -> bool {
240        // 2026-09-25: The NVFP4 and BF16 attention arms shard Q/K/V/O, and the GDN dequant
241        // path shards by head (`shard_gdn_*`). The native FP8 attention and FFN arms, the
242        // UInt8 attention weights and the keep-packed Q2 arms require TP=1; the native FP8
243        // and pre-quantized NVFP4 GDN arms do not shard.
244        true
245    }
246
247    fn load_layers(
248        &self,
249        store: &WeightStore,
250        config: &ModelConfig,
251        gpu: &dyn GpuBackend,
252        layer_kv_dtypes: &[KvCacheDtype],
253    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
254        let layer_types = if config.layer_types.is_empty() {
255            (0..config.num_hidden_layers)
256                .map(|i| config.layer_type(i))
257                .collect::<Vec<_>>()
258        } else {
259            config.layer_types.clone()
260        };
261
262        let mut layers: Vec<Box<dyn TransformerLayer>> =
263            Vec::with_capacity(config.num_hidden_layers);
264        let mut attn_idx = 0usize;
265
266        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
267        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
268        let stream = gpu.default_stream();
269        let h = config.hidden_size;
270
271        let variant = detect_nvfp4_variant(store, config);
272        let weight_format = WeightFormat::detect(store, config);
273        tracing::info!(
274            "Weight format: {:?}, NVFP4 variant: {:?}",
275            weight_format,
276            variant
277        );
278
279        // 2026-09-25: Unless `METRALE_NO_GDN_FP8_PREFILL` is set, the GDN dequant path below
280        // also builds FP8 casts of its BF16 `[Q|K|V|Z]` and out_proj (`bf16_to_fp8`).
281        let fp8_ssm_prefill = std::env::var_os("METRALE_NO_GDN_FP8_PREFILL").is_none();
282        let bf16_to_fp8_k = if fp8_ssm_prefill {
283            tracing::info!(
284                "SSM in_proj_qkv + out_proj via native FP8 prefill GEMM \
285                 (BF16 act × FP8 weight via fp8_gemm_n128). PREFILL ONLY: \
286                 decode + batched verify read the NVFP4 copy — weight-streaming \
287                 GEMV at M<=8, tile GEMM above it (trait_decode_batched.rs). \
288                 The FP8 copy reaches decode only at M<=8 on a build whose \
289                 batched NVFP4 GEMVs are absent"
290            );
291            Some(gpu.kernel("w4a16", "bf16_to_fp8")?)
292        } else {
293            None
294        };
295
296        // 2026-09-25: `METRALE_MEM_PROFILE` (any value) logs free GPU memory at the start
297        // and at every 8th layer.
298        let mem_profile = std::env::var("METRALE_MEM_PROFILE").is_ok();
299        let log_free = |tag: &str| {
300            if mem_profile && let Ok(free) = gpu.free_memory() {
301                tracing::info!("MEM_PROFILE[{tag}]: {:.2} GB GPU-free", free as f64 / 1e9);
302            }
303        };
304        log_free("dense-load-start");
305
306        // 2026-09-25: The routing levers are read once per load (`RouteEnv::from_env`,
307        // `fp8_residency.rs`).
308        let route_env = RouteEnv::from_env();
309        let mut residency = DerivedResidency::default();
310        let policy = load_cx::weight_quant_policy(config);
311        let served = std::cell::RefCell::new(served_formats::ServedFormats::default());
312        let cx = LoadCx {
313            store,
314            config,
315            gpu,
316            layer_kv_dtypes,
317            variant,
318            absmax_k,
319            quantize_k,
320            stream,
321            h,
322            bf16_to_fp8_k,
323            route_env: &route_env,
324            policy,
325            served: &served,
326        };
327
328        for (i, lt) in layer_types.iter().enumerate() {
329            if i % 8 == 0 {
330                log_free(&format!("layer-{i}"));
331            }
332            let lp = config.layer_prefix(i);
333            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
334            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;
335            let ffn = ffn_arm::build_dense_ffn(&cx, &mut residency, i, &lp)?;
336
337            match lt {
338                LayerType::FullAttention => {
339                    let l = LayerIn {
340                        i,
341                        lp: &lp,
342                        input_norm,
343                        post_attn_norm,
344                        ffn,
345                    };
346                    if let Flow::Continue = attn_layer::load_full_attention(
347                        &cx,
348                        &mut residency,
349                        &mut layers,
350                        &mut attn_idx,
351                        l,
352                    )? {
353                        continue;
354                    }
355                }
356                LayerType::LinearAttention => {
357                    let l = LayerIn {
358                        i,
359                        lp: &lp,
360                        input_norm,
361                        post_attn_norm,
362                        ffn,
363                    };
364                    if let Flow::Continue =
365                        gdn_layer::load_linear_attention(&cx, &mut residency, &mut layers, l)?
366                    {
367                        continue;
368                    }
369                }
370                LayerType::SlidingAttention => {
371                    unreachable!("unexpected SlidingAttention in this loader")
372                }
373                LayerType::Moe => unreachable!("Qwen3.5 dense has no standalone MoE layers"),
374                // 2026-09-25: GLM-5.3 `deepseek_sparse_attention` needs a DSA indexer and per-query
375                // top-k, which this loader does not build; it is refused, not served as dense
376                // attention.
377                LayerType::SparseAttention => anyhow::bail!(
378                    "layer {i}: SparseAttention needs a DSA indexer and per-query top-k; Qwen3.5 dense has neither"
379                ),
380            }
381
382            if (i + 1) % 10 == 0 {
383                tracing::info!("Loaded layers 0..{}", i + 1);
384                metrale_telemetry::progress::layer(i + 1, config.num_hidden_layers);
385            }
386        }
387
388        cx.install_w8a8_and_log_served(&layer_types, &mut layers)?;
389        tracing::info!(
390            "Qwen3.5 dense weight loader: {} layers ({} attention, {} SSM, dense FFN)",
391            layers.len(),
392            attn_idx,
393            layers.len() - attn_idx,
394        );
395        // 2026-09-25: Logged at info on every load: `weights` is the checkpoint as the store
396        // holds it, `derived` what this loader built and the store owns, `not built` what
397        // the residency plan skipped.
398        tracing::info!("{}", residency.summary(store.resident_bytes()));
399        for (label, bytes, count) in store.derived().by_label() {
400            tracing::debug!(
401                "  derived-weight owner: {:>9.1} MB x{:<5} {label}",
402                bytes as f64 / (1024.0 * 1024.0),
403                count,
404            );
405        }
406
407        Ok(layers)
408    }
409
410    /// 2026-09-25: Frees the store tensors that `load_layers` copied into fused buffers and
411    /// no longer reads: `in_proj_qkv`, `in_proj_z` (with their scales), `in_proj_a` and
412    /// `in_proj_b` where `gdn_fp8_arm_selected` holds, and `mlp.gate_proj`/`mlp.up_proj`
413    /// (with their scales) where `ffn_gateup_fused_selected` holds. `out_proj`, `conv1d`,
414    /// `A_log`, `dt_bias`, `norm` and `mlp.down_proj` are kept: layers may still alias them.
415    fn prune_after_load(
416        &self,
417        store: &mut WeightStore,
418        config: &ModelConfig,
419        gpu: &dyn GpuBackend,
420    ) -> Result<()> {
421        let doomed = prune::doomed_after_load(store, config);
422        if doomed.is_empty() {
423            return Ok(());
424        }
425        let (count, bytes) = store.free_matching(gpu, |name| doomed.contains(name))?;
426        tracing::info!(
427            "native FP8: released {count} store tensors ({:.2} GB) consumed by the fused              [QKV|Z] SSM concat, the BA interleave and the dense-FFN gate+up fusion;              out_proj/conv1d/A_log/dt_bias/norm/down_proj kept (still aliased)",
428            bytes as f64 / 1e9,
429        );
430        Ok(())
431    }
432
433    fn load_embedding(
434        &self,
435        store: &WeightStore,
436        config: &ModelConfig,
437        _gpu: &dyn GpuBackend,
438    ) -> Result<DenseWeight> {
439        loaders_b::load_embedding(store, config)
440    }
441
442    fn load_final_norm(
443        &self,
444        store: &WeightStore,
445        config: &ModelConfig,
446        _gpu: &dyn GpuBackend,
447    ) -> Result<DenseWeight> {
448        loaders_b::load_final_norm(store, config)
449    }
450
451    fn load_lm_head(
452        &self,
453        store: &WeightStore,
454        config: &ModelConfig,
455        gpu: &dyn GpuBackend,
456    ) -> Result<DenseWeight> {
457        loaders_b::load_lm_head(store, config, gpu)
458    }
459
460    fn load_mtp_weights(
461        &self,
462        store: &WeightStore,
463        config: &ModelConfig,
464        gpu: &dyn GpuBackend,
465    ) -> Result<Option<MtpWeights>> {
466        if !store.contains("mtp.fc.weight") {
467            return Ok(None);
468        }
469        let variant = detect_nvfp4_variant(store, config);
470        tracing::info!(
471            "Loading dense MTP weights (variant={:?}, hidden={}, inter={})",
472            variant,
473            config.hidden_size,
474            config.intermediate_size,
475        );
476        // 2026-09-25: `load_mtp` returns a dense head (`dense_ffn` set) when the MTP layer has
477        // `mlp.gate_proj.weight` and no router.
478        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
479        if mtp.dense_ffn.is_some() {
480            tracing::info!("Dense MTP head ready (FP8 e4m3 projections + dense gate/up/down MLP)");
481        } else {
482            tracing::info!(
483                "MoE MTP head ready ({} experts) — dense loader sees MoE bundle",
484                mtp.experts.len(),
485            );
486        }
487        Ok(Some(mtp))
488    }
489
490    fn load_vision_encoder(
491        &self,
492        store: &WeightStore,
493        config: &ModelConfig,
494        gpu: &dyn GpuBackend,
495    ) -> Result<Option<metrale_model_layers::layers::VisionTower>> {
496        // 2026-09-25: The MoE loader's `load_vision_encoder` reads only `store`, `config` and
497        // `gpu`, so the dense loader reuses it.
498        super::qwen35::Qwen35WeightLoader.load_vision_encoder(store, config, gpu)
499    }
500}
metrale_model_arch/weight_loader/qwen35_dense.rs

metrale_model_arch/weight_loader/
qwen35_dense.rs
