Stream Qwen PLE embeddings

src/engine/qwen.rs

@@ -23,6 +23,7 @@ const MTP_BYTES: u64 = 1_672_575_532;
 const KV_BYTES_PER_TOKEN: u64 = 24_576;
 const GDN_STATE_BYTES: u64 = 113_246_208;
 const GDN_CONV_BYTES: u64 = 2_211_840;
+const PLE_CONV_BYTES: u64 = 184_320;
 
 #[derive(Deserialize)]
 struct Manifest {
@@ -120,9 +121,9 @@ pub(super) struct QwenModel {
 impl QwenModel {
     pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
         let loaded = load(root, context, false)?;
-        let mut paths = loaded
-            .bindings
-            .core
+        let mut bindings = loaded.bindings.core;
+        bindings.extend(loaded.bindings.ple);
+        let mut paths = bindings
             .iter()
             .map(|binding| binding.file.clone())
             .collect::<Vec<_>>();
@@ -138,9 +139,7 @@ impl QwenModel {
             map_indices.insert(path.clone(), maps.len());
             maps.push(QwenMap { path, map });
         }
-        let tensors = loaded
-            .bindings
-            .core
+        let tensors = bindings
             .into_iter()
             .map(|binding| {
                 let tensor = QwenTensor {
@@ -178,6 +177,16 @@ impl QwenModel {
         (&self.maps[index].map, &self.maps[index].path)
     }
 
+    pub(super) fn tensor_bytes<'a>(&'a self, tensor: &QwenTensor) -> Result<&'a [u8], String> {
+        let map = &self.maps[tensor.map].map;
+        let start = usize::try_from(tensor.range.start)
+            .map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
+        let end = usize::try_from(tensor.range.end)
+            .map_err(|_| format!("{} ends beyond this platform", tensor.name))?;
+        map.get(start..end)
+            .ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name))
+    }
+
     pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
         self.identity
     }
@@ -243,6 +252,48 @@ impl QwenModel {
     pub(super) fn memory(&self) -> &MemoryPlan {
         &self.memory
     }
+
+    #[cfg(test)]
+    pub(super) fn mapped_residency(&self) -> Result<(u64, u64), String> {
+        // SAFETY: sysconf is read-only and has no pointer preconditions.
+        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
+        if page <= 0 {
+            return Err("macOS did not report its virtual-memory page size".into());
+        }
+        let page = page as usize;
+        let mut core = 0_u64;
+        let mut ple = 0_u64;
+        for item in &self.maps {
+            let mut pages = vec![0_i8; item.map.len().div_ceil(page)];
+            // SAFETY: each read-only mmap and residency vector remain valid for this call.
+            if unsafe {
+                libc::mincore(
+                    item.map.as_ptr().cast_mut().cast(),
+                    item.map.len(),
+                    pages.as_mut_ptr(),
+                )
+            } != 0
+            {
+                return Err(format!(
+                    "cannot inspect residency for {}: {}",
+                    item.path.display(),
+                    std::io::Error::last_os_error()
+                ));
+            }
+            let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
+                .min(item.map.len()) as u64;
+            if item
+                .path
+                .file_name()
+                .is_some_and(|name| name == "ngram-table.safetensors")
+            {
+                ple += bytes;
+            } else {
+                core += bytes;
+            }
+        }
+        Ok((core, ple))
+    }
 }
 
 pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
@@ -587,7 +638,7 @@ pub(super) fn memory_plan(
         .checked_mul(u64::from(context))
         .ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
     let kv_and_recurrent = kv
-        .checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES)
+        .checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES + PLE_CONV_BYTES)
         .ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
     let prefill_transient = u64::from(prefill_chunk)
         .checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
@@ -630,12 +681,12 @@ mod tests {
         assert_eq!(plan.resident_core, CORE_BYTES);
         assert_eq!(plan.mapped_ple, PLE_BYTES);
         assert_eq!(plan.optional_mtp, MTP_BYTES);
-        assert_eq!(plan.kv_and_recurrent, 6_557_908_992);
+        assert_eq!(plan.kv_and_recurrent, 6_558_093_312);
         assert_eq!(plan.prefill_transient, 27_262_976);
-        assert_eq!(plan.admission, 80_000_430_099);
+        assert_eq!(plan.admission, 80_000_614_419);
         let without_mtp = memory_plan(262_144, false, 512).unwrap();
         assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
-        assert_eq!(without_mtp.admission, 78_327_854_567);
+        assert_eq!(without_mtp.admission, 78_328_038_887);
         assert!(memory_plan(0, false, 512).is_err());
         assert!(memory_plan(262_145, false, 512).is_err());
         assert!(memory_plan(1, false, 0).is_err());