spark_model/layers/qwen3_ssm/
init_sequential.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Sequential-QKVZ constructor variant, split from `init.rs` for the
4//! ≤500 LoC cap. Child module of `init` so it shares the same `super`
5//! (the `qwen3_ssm` module) and field visibility.
6
7use anyhow::Result;
8use spark_runtime::gpu::GpuBackend;
9
10use super::super::Qwen3SsmLayer;
11use crate::layers::FfnComponent;
12use crate::weight_map::{DenseWeight, QuantizedWeight, SsmWeights};
13
14impl Qwen3SsmLayer {
15    /// Construct an SSM layer where QKVZ projection output is already sequential.
16    ///
17    /// Used by Qwen3.5 where separate QKV and Z weights are concatenated at load
18    /// time into `[Q|K|V|Z]` row order. The `deinterleave_qkvz` kernel is skipped
19    /// and plain `w4a16_gemv` writes directly to the deinterleaved buffer.
20    pub fn new_sequential(
21        input_norm: DenseWeight,
22        ssm: SsmWeights,
23        post_attn_norm: DenseWeight,
24        ffn: FfnComponent,
25        qkvz_nvfp4: Option<QuantizedWeight>,
26        qkvz_nvfp4_t: Option<QuantizedWeight>,
27        out_proj_nvfp4_t: Option<QuantizedWeight>,
28        config: &atlas_core::config::ModelConfig,
29        gpu: &dyn GpuBackend,
30    ) -> Result<Self> {
31        let mut layer = Self::new(
32            input_norm,
33            ssm,
34            post_attn_norm,
35            ffn,
36            qkvz_nvfp4,
37            config,
38            gpu,
39        )?;
40        layer.sequential_qkvz = true;
41        layer.qkvz_nvfp4_t = qkvz_nvfp4_t;
42        layer.out_proj_nvfp4_t = out_proj_nvfp4_t;
43        Ok(layer)
44    }
45}