metrale_model_layers/layers/qwen3_ssm/
init_sequential.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen3SsmLayer::new_sequential`, the constructor for
4//! checkpoints whose fused QKVZ weight is already in `[Q|K|V|Z]` row order.
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::GpuBackend;
11
12use super::super::Qwen3SsmLayer;
13use crate::layers::FfnComponent;
14use crate::weight_map::{DenseWeight, QuantizedWeight, SsmWeights};
15
16impl Qwen3SsmLayer {
17    /// 2026-09-25: `new` plus `sequential_qkvz = true` and the two transposed
18    /// NVFP4 prefill weights. With `sequential_qkvz` the QKVZ projection writes
19    /// straight into the deinterleaved buffer and `deinterleave_qkvz` is not
20    /// launched. Called by the `qwen35` and `qwen35_dense` weight loaders.
21    pub fn new_sequential(
22        input_norm: DenseWeight,
23        ssm: SsmWeights,
24        post_attn_norm: DenseWeight,
25        ffn: FfnComponent,
26        qkvz_nvfp4: Option<QuantizedWeight>,
27        qkvz_nvfp4_t: Option<QuantizedWeight>,
28        out_proj_nvfp4_t: Option<QuantizedWeight>,
29        config: &metrale_config::ModelConfig,
30        gpu: &dyn GpuBackend,
31    ) -> Result<Self> {
32        let mut layer = Self::new(
33            input_norm,
34            ssm,
35            post_attn_norm,
36            ffn,
37            qkvz_nvfp4,
38            config,
39            gpu,
40        )?;
41        layer.sequential_qkvz = true;
42        layer.qkvz_nvfp4_t = qkvz_nvfp4_t;
43        layer.out_proj_nvfp4_t = out_proj_nvfp4_t;
44        Ok(layer)
45    }
46}
metrale_model_layers/layers/qwen3_ssm/init_sequential.rs

metrale_model_layers/layers/qwen3_ssm/
init_sequential.rs
