ferrum_interfaces/vnext/
weight_source.rs

1use std::borrow::Cow;
2use std::fmt;
3use std::sync::Arc;
4
5use super::{ElementType, VNextError, WeightComponentSpec, WeightId};
6
7/// Owner for host bytes whose address, length, and contents remain stable for
8/// the lifetime of the owner.
9///
10/// # Safety
11///
12/// Implementations must return the same readable allocation from
13/// [`Self::stable_bytes`] for their entire lifetime. The allocation must not be
14/// mutated while a retained region exists. Device backends may keep a native
15/// no-copy view after the source object that produced a payload has been
16/// dropped.
17pub unsafe trait StableHostMemory: Send + Sync + 'static {
18    fn stable_bytes(&self) -> &[u8];
19}
20
21/// An owned, bounds-checked subregion of stable host memory.
22#[derive(Clone)]
23pub struct RetainedHostMemoryRegion {
24    owner: Arc<dyn StableHostMemory>,
25    offset_bytes: usize,
26    length_bytes: usize,
27}
28
29impl RetainedHostMemoryRegion {
30    pub fn new<T>(
31        owner: Arc<T>,
32        offset_bytes: usize,
33        length_bytes: usize,
34    ) -> Result<Self, VNextError>
35    where
36        T: StableHostMemory,
37    {
38        let end = offset_bytes.checked_add(length_bytes).ok_or_else(|| {
39            VNextError::InvalidExecutionPlan {
40                reason: "retained host-memory range overflows the host address space".to_owned(),
41            }
42        })?;
43        if length_bytes == 0 || end > owner.stable_bytes().len() {
44            return Err(VNextError::InvalidExecutionPlan {
45                reason: "retained host-memory range is empty or exceeds its owner".to_owned(),
46            });
47        }
48        Ok(Self {
49            owner,
50            offset_bytes,
51            length_bytes,
52        })
53    }
54
55    pub fn bytes(&self) -> &[u8] {
56        &self.owner.stable_bytes()[self.offset_bytes..self.offset_bytes + self.length_bytes]
57    }
58
59    /// Entire stable allocation. Backends use this only to prove that a
60    /// page-aligned native view enclosing [`Self::bytes`] remains within the
61    /// retained owner.
62    pub fn owner_bytes(&self) -> &[u8] {
63        self.owner.stable_bytes()
64    }
65
66    pub const fn offset_bytes(&self) -> usize {
67        self.offset_bytes
68    }
69
70    pub const fn length_bytes(&self) -> usize {
71        self.length_bytes
72    }
73}
74
75impl fmt::Debug for RetainedHostMemoryRegion {
76    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
77        formatter
78            .debug_struct("RetainedHostMemoryRegion")
79            .field("offset_bytes", &self.offset_bytes)
80            .field("length_bytes", &self.length_bytes)
81            .finish_non_exhaustive()
82    }
83}
84
85/// Validated physical bytes for one model weight component.
86///
87/// Borrowed bytes allow mmap-backed dense weights to avoid a host copy.
88/// Owned bytes cover format adapters that must repack before device upload.
89pub struct WeightComponentPayload<'source> {
90    component_id: WeightId,
91    external_names: Vec<String>,
92    source_files: Vec<String>,
93    dimensions: Vec<u64>,
94    element_type: ElementType,
95    bytes: Cow<'source, [u8]>,
96    retained_host_memory: Option<RetainedHostMemoryRegion>,
97}
98
99impl<'source> WeightComponentPayload<'source> {
100    pub fn new(
101        component: &WeightComponentSpec,
102        external_name: impl Into<String>,
103        source_file: impl Into<String>,
104        dimensions: Vec<u64>,
105        element_type: ElementType,
106        bytes: impl Into<Cow<'source, [u8]>>,
107    ) -> Result<Self, VNextError> {
108        Self::from_ordered_sources(
109            component,
110            vec![external_name.into()],
111            vec![source_file.into()],
112            dimensions,
113            element_type,
114            bytes,
115        )
116    }
117
118    /// Construct a payload materialized from multiple ordered checkpoint
119    /// tensors. Multi-source order is part of the component schema, so packed
120    /// projections cannot silently swap their logical partitions.
121    pub fn from_ordered_sources(
122        component: &WeightComponentSpec,
123        external_names: Vec<String>,
124        source_files: Vec<String>,
125        dimensions: Vec<u64>,
126        element_type: ElementType,
127        bytes: impl Into<Cow<'source, [u8]>>,
128    ) -> Result<Self, VNextError> {
129        let bytes = bytes.into();
130        let valid_source_file = |source_file: &str| {
131            !source_file.is_empty()
132                && !source_file.starts_with('/')
133                && !source_file.contains('\\')
134                && source_file
135                    .split('/')
136                    .all(|component| !matches!(component, "" | "." | ".."))
137        };
138        let sources_match = !external_names.is_empty()
139            && external_names == component.external_names
140            && external_names.len() == source_files.len()
141            && source_files.iter().all(|file| valid_source_file(file));
142        let expected_bytes = component.physical_bytes()?;
143        if !sources_match
144            || dimensions != component.dimensions
145            || element_type != component.physical_element_type()
146            || u64::try_from(bytes.len()).ok() != Some(expected_bytes)
147        {
148            return Err(VNextError::InvalidExecutionPlan {
149                reason: format!(
150                    "weight component `{}` payload differs from its schema identity, source, shape, type, or byte length",
151                    component.id
152                ),
153            });
154        }
155        Ok(Self {
156            component_id: component.id.clone(),
157            external_names,
158            source_files,
159            dimensions,
160            element_type,
161            bytes,
162            retained_host_memory: None,
163        })
164    }
165
166    /// Attach the stable owner for an otherwise borrowed payload. Pointer and
167    /// length identity are checked here so a backend cannot accidentally retain
168    /// a different mmap range than the bytes validated against the schema.
169    pub fn with_retained_host_memory(
170        mut self,
171        retained_host_memory: RetainedHostMemoryRegion,
172    ) -> Result<Self, VNextError> {
173        let retained_bytes = retained_host_memory.bytes();
174        if retained_bytes.len() != self.bytes.len()
175            || !std::ptr::eq(retained_bytes.as_ptr(), self.bytes.as_ptr())
176        {
177            return Err(VNextError::InvalidExecutionPlan {
178                reason: format!(
179                    "weight component `{}` retained host-memory region differs from its validated payload",
180                    self.component_id
181                ),
182            });
183        }
184        self.retained_host_memory = Some(retained_host_memory);
185        Ok(self)
186    }
187
188    pub fn component_id(&self) -> &WeightId {
189        &self.component_id
190    }
191
192    pub fn external_name(&self) -> &str {
193        &self.external_names[0]
194    }
195
196    pub fn source_file(&self) -> &str {
197        &self.source_files[0]
198    }
199
200    pub fn external_names(&self) -> &[String] {
201        &self.external_names
202    }
203
204    pub fn source_files(&self) -> &[String] {
205        &self.source_files
206    }
207
208    pub fn dimensions(&self) -> &[u64] {
209        &self.dimensions
210    }
211
212    pub const fn element_type(&self) -> ElementType {
213        self.element_type
214    }
215
216    pub fn bytes(&self) -> &[u8] {
217        &self.bytes
218    }
219
220    pub fn retained_host_memory(&self) -> Option<&RetainedHostMemoryRegion> {
221        self.retained_host_memory.as_ref()
222    }
223}
224
225/// Backend-neutral source of schema-addressed physical weight components.
226/// Implementations own checkpoint file-format discovery and source-payload
227/// validation. The execution plan's trusted [`super::WeightMaterializer`] owns
228/// any repacking or quantization before resource initialization performs
229/// placement and device submission.
230pub trait WeightComponentSource: Send + Sync {
231    fn component<'source>(
232        &'source self,
233        component: &WeightComponentSpec,
234    ) -> Result<WeightComponentPayload<'source>, VNextError>;
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::vnext::{BlockQuantizationSpec, WeightComponentRole, WeightEncoding};
241
242    struct StableBytes(Vec<u8>);
243
244    // SAFETY: the Vec is never mutated and owns one fixed allocation until it
245    // is dropped.
246    unsafe impl StableHostMemory for StableBytes {
247        fn stable_bytes(&self) -> &[u8] {
248            &self.0
249        }
250    }
251
252    fn packed_component() -> WeightComponentSpec {
253        WeightComponentSpec {
254            id: WeightId::new("component.test.gate_up").unwrap(),
255            role: WeightComponentRole::Values,
256            external_names: vec!["gate.weight".to_owned(), "up.weight".to_owned()],
257            dimensions: vec![2, 2, 2],
258            encoding: WeightEncoding::Dense {
259                element_type: ElementType::F16,
260            },
261            required: true,
262        }
263    }
264
265    #[test]
266    fn packed_payload_preserves_ordered_source_identity() {
267        let component = packed_component();
268        let payload = WeightComponentPayload::from_ordered_sources(
269            &component,
270            component.external_names.clone(),
271            vec![
272                "model-1.safetensors".to_owned(),
273                "model-2.safetensors".to_owned(),
274            ],
275            component.dimensions.clone(),
276            ElementType::F16,
277            vec![0_u8; 16],
278        )
279        .unwrap();
280        assert_eq!(payload.external_names(), component.external_names);
281        assert_eq!(payload.source_files().len(), 2);
282
283        let error = WeightComponentPayload::from_ordered_sources(
284            &component,
285            component.external_names.iter().rev().cloned().collect(),
286            vec![
287                "model-2.safetensors".to_owned(),
288                "model-1.safetensors".to_owned(),
289            ],
290            component.dimensions.clone(),
291            ElementType::F16,
292            vec![0_u8; 16],
293        )
294        .err()
295        .expect("source order is part of the packed component identity");
296        assert!(error.to_string().contains("differs from its schema"));
297    }
298
299    #[test]
300    fn block_quantized_payload_validates_block_abi_byte_size() {
301        let component = WeightComponentSpec {
302            id: WeightId::new("component.test.q4-k").unwrap(),
303            role: WeightComponentRole::PackedValues,
304            external_names: vec!["weight.q4_k".to_owned()],
305            dimensions: vec![2],
306            encoding: WeightEncoding::BlockQuantized(BlockQuantizationSpec {
307                format_id: "quantization.gguf.q4-k".to_owned().try_into().unwrap(),
308                logical_values_per_block: 256,
309                bytes_per_block: 144,
310            }),
311            required: true,
312        };
313        let payload = WeightComponentPayload::new(
314            &component,
315            "weight.q4_k",
316            "model.gguf",
317            vec![2],
318            ElementType::U8,
319            vec![0_u8; 288],
320        )
321        .unwrap();
322        assert_eq!(payload.bytes().len(), 288);
323
324        let error = WeightComponentPayload::new(
325            &component,
326            "weight.q4_k",
327            "model.gguf",
328            vec![2],
329            ElementType::U8,
330            vec![0_u8; 2],
331        )
332        .err()
333        .expect("block-grid element count must not be mistaken for byte length");
334        assert!(error.to_string().contains("differs from its schema"));
335    }
336
337    #[test]
338    fn retained_region_must_be_the_validated_payload_and_keeps_its_owner_alive() {
339        let component = packed_component();
340        let owner = Arc::new(StableBytes(vec![7_u8; 32]));
341        let retained = RetainedHostMemoryRegion::new(Arc::clone(&owner), 8, 16).unwrap();
342        let payload = WeightComponentPayload::from_ordered_sources(
343            &component,
344            component.external_names.clone(),
345            vec![
346                "model-1.safetensors".to_owned(),
347                "model-2.safetensors".to_owned(),
348            ],
349            component.dimensions.clone(),
350            ElementType::F16,
351            retained.bytes(),
352        )
353        .unwrap()
354        .with_retained_host_memory(retained.clone())
355        .unwrap();
356        let retained = payload.retained_host_memory().unwrap().clone();
357        drop(payload);
358        drop(owner);
359        assert_eq!(retained.bytes(), &[7_u8; 16]);
360
361        let other = Arc::new(StableBytes(vec![7_u8; 16]));
362        let wrong = RetainedHostMemoryRegion::new(other, 0, 16).unwrap();
363        let result = WeightComponentPayload::from_ordered_sources(
364            &component,
365            component.external_names.clone(),
366            vec![
367                "model-1.safetensors".to_owned(),
368                "model-2.safetensors".to_owned(),
369            ],
370            component.dimensions.clone(),
371            ElementType::F16,
372            vec![7_u8; 16],
373        )
374        .unwrap()
375        .with_retained_host_memory(wrong);
376        let error = match result {
377            Ok(_) => panic!("a different allocation must not satisfy retained payload identity"),
378            Err(error) => error,
379        };
380        assert!(error
381            .to_string()
382            .contains("retained host-memory region differs"));
383    }
384}
ferrum_interfaces/vnext/weight_source.rs

ferrum_interfaces/vnext/
weight_source.rs
