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/// One ordered byte segment of a physical weight component.
86///
87/// A segment may borrow directly from immutable checkpoint storage. When a
88/// backend needs to retain that storage beyond the source request, the
89/// attached [`RetainedHostMemoryRegion`] proves and owns the exact borrowed
90/// range.
91pub struct WeightComponentSegment<'source> {
92    bytes: Cow<'source, [u8]>,
93    retained_host_memory: Option<RetainedHostMemoryRegion>,
94}
95
96impl<'source> WeightComponentSegment<'source> {
97    pub fn new(bytes: impl Into<Cow<'source, [u8]>>) -> Self {
98        Self {
99            bytes: bytes.into(),
100            retained_host_memory: None,
101        }
102    }
103
104    /// Attach the stable owner for this segment. Pointer and length identity
105    /// are checked so a backend cannot retain a different mmap range than the
106    /// segment validated against the component schema.
107    pub fn with_retained_host_memory(
108        mut self,
109        retained_host_memory: RetainedHostMemoryRegion,
110    ) -> Result<Self, VNextError> {
111        let retained_bytes = retained_host_memory.bytes();
112        if retained_bytes.len() != self.bytes.len()
113            || !std::ptr::eq(retained_bytes.as_ptr(), self.bytes.as_ptr())
114        {
115            return Err(VNextError::InvalidExecutionPlan {
116                reason: "weight component segment retained host-memory region differs from its validated bytes"
117                    .to_owned(),
118            });
119        }
120        self.retained_host_memory = Some(retained_host_memory);
121        Ok(self)
122    }
123
124    pub fn bytes(&self) -> &[u8] {
125        &self.bytes
126    }
127
128    pub fn retained_host_memory(&self) -> Option<&RetainedHostMemoryRegion> {
129        self.retained_host_memory.as_ref()
130    }
131}
132
133/// Validated physical bytes for one model weight component split into
134/// deterministic source-order segments.
135///
136/// Segment boundaries are transport metadata, not part of the component ABI:
137/// their concatenated bytes must exactly match the component's physical byte
138/// length. This lets existing contiguous sources use one segment while
139/// adapters for packed multi-tensor components expose one mmap-backed segment
140/// per ordered checkpoint tensor without allocating an aggregate `Vec`.
141pub struct WeightComponentSegments<'source> {
142    component_id: WeightId,
143    external_names: Vec<String>,
144    source_files: Vec<String>,
145    dimensions: Vec<u64>,
146    element_type: ElementType,
147    segments: Vec<WeightComponentSegment<'source>>,
148}
149
150impl<'source> WeightComponentSegments<'source> {
151    pub fn from_ordered_segments(
152        component: &WeightComponentSpec,
153        external_names: Vec<String>,
154        source_files: Vec<String>,
155        dimensions: Vec<u64>,
156        element_type: ElementType,
157        segments: Vec<WeightComponentSegment<'source>>,
158    ) -> Result<Self, VNextError> {
159        let total_bytes = segments.iter().try_fold(0_u64, |total, segment| {
160            u64::try_from(segment.bytes().len())
161                .ok()
162                .and_then(|length| total.checked_add(length))
163        });
164        let sources_match = valid_ordered_sources(component, &external_names, &source_files);
165        let expected_bytes = component.physical_bytes()?;
166        if segments.is_empty()
167            || segments.iter().any(|segment| segment.bytes().is_empty())
168            || !sources_match
169            || dimensions != component.dimensions
170            || element_type != component.physical_element_type()
171            || total_bytes != Some(expected_bytes)
172        {
173            return Err(VNextError::InvalidExecutionPlan {
174                reason: format!(
175                    "weight component `{}` segments differ from its schema identity, source, shape, type, or byte length",
176                    component.id
177                ),
178            });
179        }
180        Ok(Self {
181            component_id: component.id.clone(),
182            external_names,
183            source_files,
184            dimensions,
185            element_type,
186            segments,
187        })
188    }
189
190    /// Wrap an already validated contiguous payload as one segment. This is
191    /// the compatibility path used by the default
192    /// [`WeightComponentSource::component_segments`] implementation.
193    pub fn from_payload(payload: WeightComponentPayload<'source>) -> Self {
194        let WeightComponentPayload {
195            component_id,
196            external_names,
197            source_files,
198            dimensions,
199            element_type,
200            bytes,
201            retained_host_memory,
202        } = payload;
203        Self {
204            component_id,
205            external_names,
206            source_files,
207            dimensions,
208            element_type,
209            segments: vec![WeightComponentSegment {
210                bytes,
211                retained_host_memory,
212            }],
213        }
214    }
215
216    pub fn component_id(&self) -> &WeightId {
217        &self.component_id
218    }
219
220    pub fn external_names(&self) -> &[String] {
221        &self.external_names
222    }
223
224    pub fn source_files(&self) -> &[String] {
225        &self.source_files
226    }
227
228    pub fn dimensions(&self) -> &[u64] {
229        &self.dimensions
230    }
231
232    pub const fn element_type(&self) -> ElementType {
233        self.element_type
234    }
235
236    pub fn segments(&self) -> &[WeightComponentSegment<'source>] {
237        &self.segments
238    }
239
240    pub fn total_bytes(&self) -> u64 {
241        self.segments
242            .iter()
243            .map(|segment| segment.bytes().len() as u64)
244            .sum()
245    }
246}
247
248/// Validated physical bytes for one model weight component.
249///
250/// Borrowed bytes allow mmap-backed dense weights to avoid a host copy.
251/// Owned bytes cover format adapters that must repack before device upload.
252pub struct WeightComponentPayload<'source> {
253    component_id: WeightId,
254    external_names: Vec<String>,
255    source_files: Vec<String>,
256    dimensions: Vec<u64>,
257    element_type: ElementType,
258    bytes: Cow<'source, [u8]>,
259    retained_host_memory: Option<RetainedHostMemoryRegion>,
260}
261
262impl<'source> WeightComponentPayload<'source> {
263    pub fn new(
264        component: &WeightComponentSpec,
265        external_name: impl Into<String>,
266        source_file: impl Into<String>,
267        dimensions: Vec<u64>,
268        element_type: ElementType,
269        bytes: impl Into<Cow<'source, [u8]>>,
270    ) -> Result<Self, VNextError> {
271        Self::from_ordered_sources(
272            component,
273            vec![external_name.into()],
274            vec![source_file.into()],
275            dimensions,
276            element_type,
277            bytes,
278        )
279    }
280
281    /// Construct a payload materialized from multiple ordered checkpoint
282    /// tensors. Multi-source order is part of the component schema, so packed
283    /// projections cannot silently swap their logical partitions.
284    pub fn from_ordered_sources(
285        component: &WeightComponentSpec,
286        external_names: Vec<String>,
287        source_files: Vec<String>,
288        dimensions: Vec<u64>,
289        element_type: ElementType,
290        bytes: impl Into<Cow<'source, [u8]>>,
291    ) -> Result<Self, VNextError> {
292        let bytes = bytes.into();
293        let sources_match = valid_ordered_sources(component, &external_names, &source_files);
294        let expected_bytes = component.physical_bytes()?;
295        if !sources_match
296            || dimensions != component.dimensions
297            || element_type != component.physical_element_type()
298            || u64::try_from(bytes.len()).ok() != Some(expected_bytes)
299        {
300            return Err(VNextError::InvalidExecutionPlan {
301                reason: format!(
302                    "weight component `{}` payload differs from its schema identity, source, shape, type, or byte length",
303                    component.id
304                ),
305            });
306        }
307        Ok(Self {
308            component_id: component.id.clone(),
309            external_names,
310            source_files,
311            dimensions,
312            element_type,
313            bytes,
314            retained_host_memory: None,
315        })
316    }
317
318    /// Attach the stable owner for an otherwise borrowed payload. Pointer and
319    /// length identity are checked here so a backend cannot accidentally retain
320    /// a different mmap range than the bytes validated against the schema.
321    pub fn with_retained_host_memory(
322        mut self,
323        retained_host_memory: RetainedHostMemoryRegion,
324    ) -> Result<Self, VNextError> {
325        let retained_bytes = retained_host_memory.bytes();
326        if retained_bytes.len() != self.bytes.len()
327            || !std::ptr::eq(retained_bytes.as_ptr(), self.bytes.as_ptr())
328        {
329            return Err(VNextError::InvalidExecutionPlan {
330                reason: format!(
331                    "weight component `{}` retained host-memory region differs from its validated payload",
332                    self.component_id
333                ),
334            });
335        }
336        self.retained_host_memory = Some(retained_host_memory);
337        Ok(self)
338    }
339
340    pub fn component_id(&self) -> &WeightId {
341        &self.component_id
342    }
343
344    pub fn external_name(&self) -> &str {
345        &self.external_names[0]
346    }
347
348    pub fn source_file(&self) -> &str {
349        &self.source_files[0]
350    }
351
352    pub fn external_names(&self) -> &[String] {
353        &self.external_names
354    }
355
356    pub fn source_files(&self) -> &[String] {
357        &self.source_files
358    }
359
360    pub fn dimensions(&self) -> &[u64] {
361        &self.dimensions
362    }
363
364    pub const fn element_type(&self) -> ElementType {
365        self.element_type
366    }
367
368    pub fn bytes(&self) -> &[u8] {
369        &self.bytes
370    }
371
372    pub fn retained_host_memory(&self) -> Option<&RetainedHostMemoryRegion> {
373        self.retained_host_memory.as_ref()
374    }
375}
376
377/// Backend-neutral source of schema-addressed physical weight components.
378/// Implementations own checkpoint file-format discovery and source-payload
379/// validation. The execution plan's trusted [`super::WeightMaterializer`] owns
380/// any repacking or quantization before resource initialization performs
381/// placement and device submission.
382pub trait WeightComponentSource: Send + Sync {
383    fn component<'source>(
384        &'source self,
385        component: &WeightComponentSpec,
386    ) -> Result<WeightComponentPayload<'source>, VNextError>;
387
388    /// Return deterministic source-order segments for one component.
389    ///
390    /// Existing sources remain compatible through a one-segment wrapper.
391    /// Format adapters may override this method to expose independently
392    /// retained mmap ranges without changing the contiguous `component()`
393    /// contract used by existing materializers and upload paths.
394    fn component_segments<'source>(
395        &'source self,
396        component: &WeightComponentSpec,
397    ) -> Result<WeightComponentSegments<'source>, VNextError> {
398        self.component(component)
399            .map(WeightComponentSegments::from_payload)
400    }
401}
402
403fn valid_ordered_sources(
404    component: &WeightComponentSpec,
405    external_names: &[String],
406    source_files: &[String],
407) -> bool {
408    let valid_source_file = |source_file: &str| {
409        !source_file.is_empty()
410            && !source_file.starts_with('/')
411            && !source_file.contains('\\')
412            && source_file
413                .split('/')
414                .all(|component| !matches!(component, "" | "." | ".."))
415    };
416    !external_names.is_empty()
417        && external_names == component.external_names
418        && external_names.len() == source_files.len()
419        && source_files.iter().all(|file| valid_source_file(file))
420}
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425    use crate::vnext::{BlockQuantizationSpec, WeightComponentRole, WeightEncoding};
426
427    struct StableBytes(Vec<u8>);
428
429    // SAFETY: the Vec is never mutated and owns one fixed allocation until it
430    // is dropped.
431    unsafe impl StableHostMemory for StableBytes {
432        fn stable_bytes(&self) -> &[u8] {
433            &self.0
434        }
435    }
436
437    struct ContiguousSource(Vec<u8>);
438
439    impl WeightComponentSource for ContiguousSource {
440        fn component<'source>(
441            &'source self,
442            component: &WeightComponentSpec,
443        ) -> Result<WeightComponentPayload<'source>, VNextError> {
444            WeightComponentPayload::from_ordered_sources(
445                component,
446                component.external_names.clone(),
447                vec![
448                    "model-1.safetensors".to_owned(),
449                    "model-2.safetensors".to_owned(),
450                ],
451                component.dimensions.clone(),
452                ElementType::F16,
453                self.0.as_slice(),
454            )
455        }
456    }
457
458    fn packed_component() -> WeightComponentSpec {
459        WeightComponentSpec {
460            id: WeightId::new("component.test.gate_up").unwrap(),
461            role: WeightComponentRole::Values,
462            external_names: vec!["gate.weight".to_owned(), "up.weight".to_owned()],
463            dimensions: vec![2, 2, 2],
464            encoding: WeightEncoding::Dense {
465                element_type: ElementType::F16,
466            },
467            required: true,
468        }
469    }
470
471    #[test]
472    fn packed_payload_preserves_ordered_source_identity() {
473        let component = packed_component();
474        let payload = WeightComponentPayload::from_ordered_sources(
475            &component,
476            component.external_names.clone(),
477            vec![
478                "model-1.safetensors".to_owned(),
479                "model-2.safetensors".to_owned(),
480            ],
481            component.dimensions.clone(),
482            ElementType::F16,
483            vec![0_u8; 16],
484        )
485        .unwrap();
486        assert_eq!(payload.external_names(), component.external_names);
487        assert_eq!(payload.source_files().len(), 2);
488
489        let error = WeightComponentPayload::from_ordered_sources(
490            &component,
491            component.external_names.iter().rev().cloned().collect(),
492            vec![
493                "model-2.safetensors".to_owned(),
494                "model-1.safetensors".to_owned(),
495            ],
496            component.dimensions.clone(),
497            ElementType::F16,
498            vec![0_u8; 16],
499        )
500        .err()
501        .expect("source order is part of the packed component identity");
502        assert!(error.to_string().contains("differs from its schema"));
503    }
504
505    #[test]
506    fn block_quantized_payload_validates_block_abi_byte_size() {
507        let component = WeightComponentSpec {
508            id: WeightId::new("component.test.q4-k").unwrap(),
509            role: WeightComponentRole::PackedValues,
510            external_names: vec!["weight.q4_k".to_owned()],
511            dimensions: vec![2],
512            encoding: WeightEncoding::BlockQuantized(BlockQuantizationSpec {
513                format_id: "quantization.gguf.q4-k".to_owned().try_into().unwrap(),
514                logical_values_per_block: 256,
515                bytes_per_block: 144,
516            }),
517            required: true,
518        };
519        let payload = WeightComponentPayload::new(
520            &component,
521            "weight.q4_k",
522            "model.gguf",
523            vec![2],
524            ElementType::U8,
525            vec![0_u8; 288],
526        )
527        .unwrap();
528        assert_eq!(payload.bytes().len(), 288);
529
530        let error = WeightComponentPayload::new(
531            &component,
532            "weight.q4_k",
533            "model.gguf",
534            vec![2],
535            ElementType::U8,
536            vec![0_u8; 2],
537        )
538        .err()
539        .expect("block-grid element count must not be mistaken for byte length");
540        assert!(error.to_string().contains("differs from its schema"));
541    }
542
543    #[test]
544    fn retained_region_must_be_the_validated_payload_and_keeps_its_owner_alive() {
545        let component = packed_component();
546        let owner = Arc::new(StableBytes(vec![7_u8; 32]));
547        let retained = RetainedHostMemoryRegion::new(Arc::clone(&owner), 8, 16).unwrap();
548        let payload = WeightComponentPayload::from_ordered_sources(
549            &component,
550            component.external_names.clone(),
551            vec![
552                "model-1.safetensors".to_owned(),
553                "model-2.safetensors".to_owned(),
554            ],
555            component.dimensions.clone(),
556            ElementType::F16,
557            retained.bytes(),
558        )
559        .unwrap()
560        .with_retained_host_memory(retained.clone())
561        .unwrap();
562        let retained = payload.retained_host_memory().unwrap().clone();
563        drop(payload);
564        drop(owner);
565        assert_eq!(retained.bytes(), &[7_u8; 16]);
566
567        let other = Arc::new(StableBytes(vec![7_u8; 16]));
568        let wrong = RetainedHostMemoryRegion::new(other, 0, 16).unwrap();
569        let result = WeightComponentPayload::from_ordered_sources(
570            &component,
571            component.external_names.clone(),
572            vec![
573                "model-1.safetensors".to_owned(),
574                "model-2.safetensors".to_owned(),
575            ],
576            component.dimensions.clone(),
577            ElementType::F16,
578            vec![7_u8; 16],
579        )
580        .unwrap()
581        .with_retained_host_memory(wrong);
582        let error = match result {
583            Ok(_) => panic!("a different allocation must not satisfy retained payload identity"),
584            Err(error) => error,
585        };
586        assert!(error
587            .to_string()
588            .contains("retained host-memory region differs"));
589    }
590
591    #[test]
592    fn ordered_segments_preserve_source_order_and_retain_each_exact_range() {
593        let component = packed_component();
594        let owner_a = Arc::new(StableBytes(vec![1_u8; 24]));
595        let owner_b = Arc::new(StableBytes(vec![2_u8; 24]));
596        let retained_a = RetainedHostMemoryRegion::new(Arc::clone(&owner_a), 4, 8).unwrap();
597        let retained_b = RetainedHostMemoryRegion::new(Arc::clone(&owner_b), 12, 8).unwrap();
598        let segments = WeightComponentSegments::from_ordered_segments(
599            &component,
600            component.external_names.clone(),
601            vec![
602                "model-1.safetensors".to_owned(),
603                "model-2.safetensors".to_owned(),
604            ],
605            component.dimensions.clone(),
606            ElementType::F16,
607            vec![
608                WeightComponentSegment::new(retained_a.bytes())
609                    .with_retained_host_memory(retained_a.clone())
610                    .unwrap(),
611                WeightComponentSegment::new(retained_b.bytes())
612                    .with_retained_host_memory(retained_b.clone())
613                    .unwrap(),
614            ],
615        )
616        .unwrap();
617
618        assert_eq!(segments.external_names(), component.external_names);
619        assert_eq!(segments.source_files().len(), 2);
620        assert_eq!(segments.total_bytes(), 16);
621        assert_eq!(segments.segments().len(), 2);
622        assert_eq!(segments.segments()[0].bytes(), &[1_u8; 8]);
623        assert_eq!(segments.segments()[1].bytes(), &[2_u8; 8]);
624        let retained_segments = segments
625            .segments()
626            .iter()
627            .map(|segment| segment.retained_host_memory().unwrap().clone())
628            .collect::<Vec<_>>();
629        drop(segments);
630        drop(retained_a);
631        drop(retained_b);
632        drop(owner_a);
633        drop(owner_b);
634        assert_eq!(retained_segments[0].bytes(), &[1_u8; 8]);
635        assert_eq!(retained_segments[1].bytes(), &[2_u8; 8]);
636    }
637
638    #[test]
639    fn component_segments_default_wraps_the_contiguous_payload_once() {
640        let component = packed_component();
641        let source = ContiguousSource(vec![3_u8; 16]);
642        let segments = source.component_segments(&component).unwrap();
643
644        assert_eq!(segments.component_id(), &component.id);
645        assert_eq!(segments.dimensions(), component.dimensions);
646        assert_eq!(segments.element_type(), ElementType::F16);
647        assert_eq!(segments.segments().len(), 1);
648        assert_eq!(segments.segments()[0].bytes(), &[3_u8; 16]);
649    }
650
651    #[test]
652    fn ordered_segments_reject_wrong_total_length_or_retained_range() {
653        let component = packed_component();
654        let error = WeightComponentSegments::from_ordered_segments(
655            &component,
656            component.external_names.clone(),
657            vec![
658                "model-1.safetensors".to_owned(),
659                "model-2.safetensors".to_owned(),
660            ],
661            component.dimensions.clone(),
662            ElementType::F16,
663            vec![WeightComponentSegment::new(vec![0_u8; 15])],
664        )
665        .err()
666        .expect("the segment total must match the physical component size");
667        assert!(error
668            .to_string()
669            .contains("segments differ from its schema"));
670
671        let bytes = vec![5_u8; 8];
672        let owner = Arc::new(StableBytes(vec![5_u8; 8]));
673        let retained = RetainedHostMemoryRegion::new(owner, 0, 8).unwrap();
674        let error = WeightComponentSegment::new(bytes)
675            .with_retained_host_memory(retained)
676            .err()
677            .expect("a different allocation must not satisfy segment retention");
678        assert!(error
679            .to_string()
680            .contains("retained host-memory region differs"));
681    }
682}
ferrum_interfaces/vnext/weight_source.rs

ferrum_interfaces/vnext/
weight_source.rs
