MCPcopy Create free account
hub / github.com/evilsocket/cake / new

Method new

cake-core/src/models/common/disk_expert_provider.rs:120–185  ·  view source on GitHub ↗

Create a disk-backed expert provider. - `storage`: safetensors storage with pread access - `layer_prefix`: e.g., "model.layers.5.mlp" — experts are at `"{prefix}.experts.{idx}.{gate_proj,up_proj,down_proj}.weight"` - `num_experts`: total number of experts in this layer - `device`: target device (CPU for inference, or GPU for transfer) - `dtype`: target dtype for the loaded tensors - `gptq_group_s

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

Source from the content-addressed store, hash-verified

118 /// - `dtype`: target dtype for the loaded tensors
119 /// - `gptq_group_size`: if Some, experts are GPTQ-quantized and will be dequantized on read
120 pub fn new(
121 storage: Arc<dyn TensorStorageProvider>,
122 layer_prefix: String,
123 num_experts: usize,
124 device: Device,
125 dtype: DType,
126 gptq_group_size: Option<usize>,
127 ) -> Self {
128 let expert_names: Vec<ExpertNames> = (0..num_experts)
129 .map(|idx| {
130 let prefix = format!("{}.experts.{}", layer_prefix, idx);
131 ExpertNames {
132 gate_proj: format!("{prefix}.gate_proj.weight"),
133 up_proj: format!("{prefix}.up_proj.weight"),
134 down_proj: format!("{prefix}.down_proj.weight"),
135 }
136 })
137 .collect();
138 let needs_device_transfer = !device.is_cpu();
139 // Auto-detect GPTQ: if caller says GPTQ, verify first expert has qweight.
140 let gptq_group_size = if let Some(gs) = gptq_group_size {
141 let qw_name = format!("{}.experts.0.gate_proj.qweight", layer_prefix);
142 if storage.has_tensor(&qw_name) {
143 log::info!("expert offload: GPTQ 4-bit detected (group_size={gs})");
144 Some(gs)
145 } else {
146 log::info!("expert offload: GPTQ requested but no qweight found, using plain weights");
147 None
148 }
149 } else {
150 None
151 };
152 // Detect storage dtype from first expert (skip for GPTQ — qweight is int32, not weights)
153 let storage_dtype = if gptq_group_size.is_none() {
154 expert_names.first().and_then(|names| {
155 // Use metadata-only lookup (no data read) → tensor_bytes fallback → full read
156 if let Some((dt, _)) = storage.tensor_meta(&names.gate_proj) {
157 Some(dt)
158 } else if let Some((_, dt, _)) = storage.tensor_bytes(&names.gate_proj) {
159 Some(dt)
160 } else {
161 storage.read_tensor(&names.gate_proj).ok().map(|d| d.dtype)
162 }
163 })
164 } else {
165 None
166 };
167 // For GPTQ experts, disable F32 zerocopy (qweight is int32, not F32)
168 let use_f32_zerocopy = gptq_group_size.is_none()
169 && dtype == DType::F32
170 && storage_dtype.is_some_and(|sd| sd == DType::F32);
171 Self {
172 storage,
173 layer_prefix,
174 expert_names,
175 num_experts,
176 device,
177 dtype,

Callers

nothing calls this directly

Calls 4

has_tensorMethod · 0.45
tensor_metaMethod · 0.45
tensor_bytesMethod · 0.45
read_tensorMethod · 0.45

Tested by

no test coverage detected