onnx_runtime_ep_api/
weight.rs

1//! Capability-negotiated lazy weight handles for executor-to-EP delivery.
2
3use std::collections::BTreeSet;
4use std::sync::Arc;
5
6use onnx_runtime_ir::DataType;
7
8use crate::ExternalMmapRegion;
9
10/// Capability flag advertised by paging-aware execution providers.
11pub const NXRT_WEIGHT_PAGING_CAPABILITY: &str = "nxrt";
12
13#[derive(Clone, Debug, Default, PartialEq, Eq)]
14pub struct ExecutionProviderCapabilities {
15    flags: BTreeSet<String>,
16}
17
18impl ExecutionProviderCapabilities {
19    pub fn stock() -> Self {
20        Self::default()
21    }
22
23    pub fn nxrt_weight_paging() -> Self {
24        Self::from_flags([NXRT_WEIGHT_PAGING_CAPABILITY])
25    }
26
27    pub fn from_flags(flags: impl IntoIterator<Item = impl Into<String>>) -> Self {
28        Self {
29            flags: flags.into_iter().map(Into::into).collect(),
30        }
31    }
32
33    pub fn advertises(&self, capability: &str) -> bool {
34        self.flags.contains(capability)
35    }
36}
37
38#[derive(Clone, Debug, PartialEq, Eq)]
39pub struct ResidentWeight {
40    pub dtype: DataType,
41    pub shape: Vec<usize>,
42    bytes: Arc<[u8]>,
43}
44
45impl ResidentWeight {
46    pub fn new(
47        dtype: DataType,
48        shape: Vec<usize>,
49        bytes: impl Into<Arc<[u8]>>,
50    ) -> Result<Self, WeightHandleError> {
51        let elements = checked_shape_product(&shape)?;
52        let expected = dtype.checked_storage_bytes(elements).ok_or_else(|| {
53            WeightHandleError::InvalidResident("resident weight byte count overflow".into())
54        })?;
55        if expected > isize::MAX as usize {
56            return Err(WeightHandleError::InvalidResident(
57                "resident weight byte count exceeds isize::MAX".into(),
58            ));
59        }
60        let bytes = bytes.into();
61        if bytes.len() != expected {
62            return Err(WeightHandleError::InvalidResident(format!(
63                "resident weight has {} bytes, expected {expected}",
64                bytes.len()
65            )));
66        }
67        Ok(Self {
68            dtype,
69            shape,
70            bytes,
71        })
72    }
73
74    pub fn bytes(&self) -> &[u8] {
75        &self.bytes
76    }
77}
78
79fn checked_shape_product(shape: &[usize]) -> Result<usize, WeightHandleError> {
80    let mut product = 1usize;
81    let mut has_zero = false;
82    for &dimension in shape {
83        if dimension == 0 {
84            has_zero = true;
85        } else {
86            product = product.checked_mul(dimension).ok_or_else(|| {
87                WeightHandleError::InvalidResident("resident weight element count overflow".into())
88            })?;
89        }
90    }
91    Ok(if has_zero { 0 } else { product })
92}
93
94#[derive(Clone, Copy, Debug, PartialEq, Eq)]
95pub enum LazyWeightBoundary {
96    /// `pkg.nxrt::BlockQuantizedMoE`, the Phase-3 offload binding boundary.
97    BlockQuantizedMoe,
98}
99
100impl LazyWeightBoundary {
101    pub fn matches(self, domain: &str, op_type: &str) -> bool {
102        matches!(self, Self::BlockQuantizedMoe)
103            && domain == "pkg.nxrt"
104            && op_type == "BlockQuantizedMoE"
105    }
106}
107
108pub trait ResidentWeightMaterializer: Send + Sync {
109    fn materialize(&self) -> Result<ResidentWeight, WeightHandleError>;
110}
111
112impl<F> ResidentWeightMaterializer for F
113where
114    F: Fn() -> Result<ResidentWeight, WeightHandleError> + Send + Sync,
115{
116    fn materialize(&self) -> Result<ResidentWeight, WeightHandleError> {
117        self()
118    }
119}
120
121#[derive(Clone)]
122pub struct LazyWeight {
123    pub boundary: LazyWeightBoundary,
124    /// Validated external mmap ranges that back this initializer.
125    pub regions: Vec<ExternalMmapRegion>,
126    resident_materializer: Arc<dyn ResidentWeightMaterializer>,
127}
128
129impl std::fmt::Debug for LazyWeight {
130    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        formatter
132            .debug_struct("LazyWeight")
133            .field("boundary", &self.boundary)
134            .field("regions", &self.regions)
135            .field("resident_materializer", &"<deferred>")
136            .finish()
137    }
138}
139
140impl LazyWeight {
141    pub fn block_quantized_moe<M>(
142        regions: Vec<ExternalMmapRegion>,
143        resident_materializer: M,
144    ) -> Result<Self, WeightHandleError>
145    where
146        M: ResidentWeightMaterializer + 'static,
147    {
148        if regions.is_empty() {
149            return Err(WeightHandleError::MissingRegions);
150        }
151        Ok(Self {
152            boundary: LazyWeightBoundary::BlockQuantizedMoe,
153            regions,
154            resident_materializer: Arc::new(resident_materializer),
155        })
156    }
157
158    /// Materialize the unchanged stock-EP resident behavior.
159    pub fn materialize(&self) -> Result<ResidentWeight, WeightHandleError> {
160        self.resident_materializer.materialize()
161    }
162}
163
164/// General executor weight input: resident today, lazy when an EP opts in.
165#[derive(Clone, Debug)]
166pub enum WeightHandle {
167    Resident(ResidentWeight),
168    Lazy(LazyWeight),
169}
170
171impl WeightHandle {
172    pub fn negotiate(
173        &self,
174        capabilities: &ExecutionProviderCapabilities,
175    ) -> Result<NegotiatedWeight, WeightHandleError> {
176        match self {
177            Self::Resident(weight) => Ok(NegotiatedWeight::Resident(weight.clone())),
178            Self::Lazy(weight) if capabilities.advertises(NXRT_WEIGHT_PAGING_CAPABILITY) => {
179                Ok(NegotiatedWeight::Lazy(weight.clone()))
180            }
181            Self::Lazy(weight) => Ok(NegotiatedWeight::Resident(weight.materialize()?)),
182        }
183    }
184
185    pub fn is_lazy_for(&self, capabilities: &ExecutionProviderCapabilities) -> bool {
186        matches!(self, Self::Lazy(_)) && capabilities.advertises(NXRT_WEIGHT_PAGING_CAPABILITY)
187    }
188}
189
190#[derive(Clone, Debug)]
191pub enum NegotiatedWeight {
192    Resident(ResidentWeight),
193    Lazy(LazyWeight),
194}
195
196impl NegotiatedWeight {
197    /// Explicit host route available to every executor and EP.
198    pub fn materialize_host_fallback(&self) -> Result<ResidentWeight, WeightHandleError> {
199        match self {
200            Self::Resident(weight) => Ok(weight.clone()),
201            Self::Lazy(weight) => weight.materialize(),
202        }
203    }
204
205    /// Phase 3b will call this at `pkg.nxrt::BlockQuantizedMoE` binding time.
206    pub fn try_bind_device<B: LazyDeviceWeightBinder>(
207        &self,
208        binder: &B,
209    ) -> Result<B::Binding, WeightHandleError> {
210        match self {
211            Self::Resident(_) => Err(WeightHandleError::Unsupported(
212                "resident weights do not require lazy device binding".into(),
213            )),
214            Self::Lazy(weight) => binder.bind_block_quantized_moe(weight),
215        }
216    }
217}
218
219/// EP seam for Phase-3b live device paging.
220pub trait LazyDeviceWeightBinder {
221    type Binding;
222
223    fn bind_block_quantized_moe(
224        &self,
225        weight: &LazyWeight,
226    ) -> Result<Self::Binding, WeightHandleError>;
227}
228
229/// CPU-only Phase-3a binder: callers must use the host materialization route.
230#[derive(Clone, Copy, Debug, Default)]
231pub struct Phase3aHostOnlyBinder;
232
233impl LazyDeviceWeightBinder for Phase3aHostOnlyBinder {
234    type Binding = ();
235
236    fn bind_block_quantized_moe(
237        &self,
238        _weight: &LazyWeight,
239    ) -> Result<Self::Binding, WeightHandleError> {
240        Err(WeightHandleError::Unsupported(
241            "live device weight paging is deferred to WEIGHT_OFFLOAD Phase 3b".into(),
242        ))
243    }
244}
245
246#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
247pub enum WeightHandleError {
248    #[error("invalid resident weight: {0}")]
249    InvalidResident(String),
250    #[error("lazy weight requires at least one external mmap region")]
251    MissingRegions,
252    #[error("unsupported: {0}")]
253    Unsupported(String),
254}
255
256#[cfg(test)]
257mod tests {
258    use std::sync::atomic::{AtomicUsize, Ordering};
259
260    use super::*;
261
262    fn resident() -> ResidentWeight {
263        ResidentWeight::new(DataType::Uint8, vec![4], vec![1, 2, 3, 4]).unwrap()
264    }
265
266    fn region() -> ExternalMmapRegion {
267        ExternalMmapRegion {
268            mapping_id: 7,
269            offset: 100,
270            len: 4,
271        }
272    }
273
274    fn lazy() -> WeightHandle {
275        WeightHandle::Lazy(
276            LazyWeight::block_quantized_moe(vec![region()], || Ok(resident())).unwrap(),
277        )
278    }
279
280    #[test]
281    fn stock_ep_materializes_the_resident_fallback() {
282        let NegotiatedWeight::Resident(weight) = lazy()
283            .negotiate(&ExecutionProviderCapabilities::stock())
284            .unwrap()
285        else {
286            panic!("stock EP must receive resident materialization");
287        };
288        assert_eq!(weight.bytes(), &[1, 2, 3, 4]);
289    }
290
291    #[test]
292    fn nxrt_capability_preserves_lazy_block_quantized_moe_handle() {
293        let materializations = Arc::new(AtomicUsize::new(0));
294        let counter = Arc::clone(&materializations);
295        let lazy = WeightHandle::Lazy(
296            LazyWeight::block_quantized_moe(vec![region()], move || {
297                counter.fetch_add(1, Ordering::Relaxed);
298                Ok(resident())
299            })
300            .unwrap(),
301        );
302        let NegotiatedWeight::Lazy(weight) = lazy
303            .negotiate(&ExecutionProviderCapabilities::nxrt_weight_paging())
304            .unwrap()
305        else {
306            panic!("nxrt EP must receive lazy weight handle");
307        };
308        assert_eq!(weight.boundary, LazyWeightBoundary::BlockQuantizedMoe);
309        assert_eq!(weight.regions, vec![region()]);
310        assert_eq!(materializations.load(Ordering::Relaxed), 0);
311        assert_eq!(weight.materialize().unwrap().bytes(), &[1, 2, 3, 4]);
312        assert_eq!(materializations.load(Ordering::Relaxed), 1);
313    }
314
315    #[test]
316    fn phase3a_device_binding_is_explicitly_unsupported_with_host_route() {
317        let negotiated = lazy()
318            .negotiate(&ExecutionProviderCapabilities::nxrt_weight_paging())
319            .unwrap();
320        assert_eq!(
321            negotiated.try_bind_device(&Phase3aHostOnlyBinder),
322            Err(WeightHandleError::Unsupported(
323                "live device weight paging is deferred to WEIGHT_OFFLOAD Phase 3b".into()
324            ))
325        );
326        assert_eq!(
327            negotiated.materialize_host_fallback().unwrap().bytes(),
328            &[1, 2, 3, 4]
329        );
330    }
331}
onnx_runtime_ep_api/weight.rs

onnx_runtime_ep_api/
weight.rs
