Method new_stacked

cake-core/src/models/common/disk_expert_provider.rs:192–330

Create a disk-backed expert provider for the stacked switch_mlp format. In this format, expert weights are stored as 3D stacked tensors: `"{prefix}.gate_proj.weight"` with shape `(num_experts, intermediate, hidden)`. Individual expert slices are read via `tensor_slice_bytes()`.

(
        storage: Arc<dyn TensorStorageProvider>,
        layer_prefix: String,
        num_experts: usize,
        device: Device,
        dtype: DType,
    )

    /// `"{prefix}.gate_proj.weight"` with shape `(num_experts, intermediate, hidden)`.
    /// Individual expert slices are read via `tensor_slice_bytes()`.
    pub fn new_stacked(
        storage: Arc<dyn TensorStorageProvider>,
        layer_prefix: String,
        num_experts: usize,
        device: Device,
        dtype: DType,
    ) -> Self {
        let expert_names: Vec<ExpertNames> = (0..num_experts)
            .map(|_| {
                ExpertNames {
                    gate_proj: format!("{layer_prefix}.gate_proj.weight"),
                    up_proj: format!("{layer_prefix}.up_proj.weight"),
                    down_proj: format!("{layer_prefix}.down_proj.weight"),
                }
            })
            .collect();
        let needs_device_transfer = !device.is_cpu();
        // Pre-compute stacked metadata from the 3D tensor shapes
        let stacked_meta = {
            let gate_meta = storage.tensor_meta(&expert_names[0].gate_proj);
            let up_meta = storage.tensor_meta(&expert_names[0].up_proj);
            let down_meta = storage.tensor_meta(&expert_names[0].down_proj);
            // Check for affine 4-bit quantization (scales + biases companions)
            let scales_name = expert_names[0].gate_proj.replace(".weight", ".scales");
            let is_affine = storage.has_tensor(&scales_name);
            match (gate_meta, up_meta, down_meta) {
                (Some((gdt, gs)), Some((_, us)), Some((_, ds))) if gs.len() == 3 && us.len() == 3 && ds.len() == 3 => {
                    let dtype_size = gdt.size_in_bytes();
                    let make_proj = |wshape: &[usize], proj_name: &str| -> StackedProjMeta {
                        let w_stride = wshape[1] * wshape[2] * dtype_size;
                        let (s_stride, s_shape) = if is_affine {
                            let sn = proj_name.replace(".weight", ".scales");
                            if let Some((sdt, ss)) = storage.tensor_meta(&sn) {
                                let ss_size = sdt.size_in_bytes();
                                (ss[1] * ss[2] * ss_size, [ss[1], ss[2]])
                            } else {
                                (0, [0, 0])
                            }
                        } else {
                            (0, [0, 0])
                        };
                        let b_stride = if is_affine {
                            let bn = proj_name.replace(".weight", ".biases");
                            if let Some((bdt, bs)) = storage.tensor_meta(&bn) {
                                bs[1] * bs[2] * bdt.size_in_bytes()
                            } else {
                                0
                            }
                        } else {
                            0
                        };
                        StackedProjMeta {
                            weight_byte_stride: w_stride,
                            weight_shape: [wshape[1], wshape[2]],
                            scales_byte_stride: s_stride,
                            scales_shape: s_shape,
                            biases_byte_stride: b_stride,
                        }