Optimize Qwen3.8 inference on Apple silicon

src/engine/metal.rs

@@ -13,10 +13,13 @@ use profile::ExpertProfile;
 use qwen::QwenExecutor;
 
 use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
+#[cfg(test)]
+use super::sample_from_logits;
 use super::validation::{DsparkConfig, SupportKind, dspark_config};
 use super::{
-    LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, exact_speculative_sample,
-    sample_from_logits,
+    LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample,
+    exact_speculative_sample_from_probabilities, qwen_sampling_probabilities,
+    sample_probabilities_f64,
 };
 use crate::model::ModelChoice;
 use crate::settings::{
@@ -4517,6 +4520,20 @@ impl Executor {
         }
     }
 
+    pub(super) fn take_pending_sampled_token(&mut self) -> Option<i32> {
+        match self {
+            Self::Qwen(executor) => executor.take_pending_sampled_token(),
+            Self::DeepSeek(_) | Self::Glm(_) => None,
+        }
+    }
+
+    pub(super) fn finalize_pending_sampled_token(&mut self) -> Result<(), String> {
+        match self {
+            Self::Qwen(executor) => executor.finalize_pending_sampled_token(),
+            Self::DeepSeek(_) | Self::Glm(_) => Ok(()),
+        }
+    }
+
     pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
         match self {
             Self::DeepSeek(executor) => executor.eval(token),

src/engine/metal/gpu.rs

@@ -8,6 +8,7 @@ pub(super) struct GpuTensor {
 #[derive(Clone, Copy)]
 #[repr(C)]
 pub(super) struct QwenWeightView {
+    pub(super) tensor: *const GpuTensor,
     pub(super) map: *const c_void,
     pub(super) size: u64,
     pub(super) offset: u64,
@@ -90,6 +91,7 @@ pub(super) struct StreamExpertCacheStats {
 unsafe extern "C" {
     pub(super) fn ds4_gpu_init() -> i32;
     pub(super) fn ds4_gpu_cleanup();
+    pub(super) fn ds4_gpu_metal4_tensor_api_enabled() -> i32;
     pub(super) fn ds4_gpu_set_model_map_range(
         model_map: *const c_void,
         model_size: u64,
@@ -103,6 +105,7 @@ unsafe extern "C" {
         a: *const GpuTensor,
         b: *const GpuTensor,
         c: *const GpuTensor,
+        d: *const GpuTensor,
         weights: *const QwenWeightView,
         weight_count: u32,
         args: *const QwenKernelArgs,
@@ -132,6 +135,7 @@ unsafe extern "C" {
     ) -> i32;
     pub(super) fn ds4_gpu_set_quality(quality: bool);
     pub(super) fn ds4_gpu_set_glm_model(enabled: bool);
+    pub(super) fn ds4_gpu_set_qwen_model(enabled: bool);
     pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool);
     pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
     pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
@@ -199,6 +203,7 @@ unsafe extern "C" {
     pub(super) fn ds4_gpu_parallel_ffn_finish() -> i32;
     pub(super) fn ds4_gpu_parallel_ffn_abort();
     pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
+    pub(super) fn ds4_gpu_tensor_alloc_untracked(bytes: u64) -> *mut GpuTensor;
     pub(super) fn ds4_gpu_tensor_view(
         base: *const GpuTensor,
         offset: u64,
@@ -1847,6 +1852,7 @@ impl Context {
         check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
         unsafe {
             ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
+            ds4_gpu_set_qwen_model(false);
             ds4_gpu_set_ssd_streaming(ssd_streaming);
             // Decode enables this only around DS4's eligible resident pre-M5
             // MXFP4 token path; all other work starts from the portable path.
@@ -1944,6 +1950,7 @@ impl Context {
         check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
         unsafe {
             ds4_gpu_set_glm_model(false);
+            ds4_gpu_set_qwen_model(true);
             ds4_gpu_set_ssd_streaming(false);
             ds4_gpu_set_decode_pipeline_fast_lookup(0);
         }
@@ -2071,6 +2078,12 @@ impl Buffer {
             .ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
     }
 
+    pub(super) fn untracked(bytes: u64) -> Result<Self, String> {
+        NonNull::new(unsafe { ds4_gpu_tensor_alloc_untracked(bytes) })
+            .map(Self)
+            .ok_or_else(|| format!("Metal could not allocate {bytes} untracked bytes"))
+    }
+
     pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
         // SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
         // the view, ARC-retains base_obj.buffer, and marks the view non-owning, so

src/engine/qwen.rs

@@ -3,15 +3,19 @@ use super::{ChatTurn, ModelSummary};
 use crate::model::ModelChoice;
 use crate::settings::ReasoningMode;
 use memmap2::{Mmap, MmapOptions};
+use rustc_hash::FxHashMap as HashMap;
 use serde::de::{MapAccess, Visitor};
 use serde::{Deserialize, Deserializer};
 use serde_json::Value;
 use sha2::{Digest, Sha256};
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque};
 use std::fmt;
 use std::fs::{self, File};
 use std::io::Read;
+use std::os::unix::fs::FileExt;
 use std::path::{Path, PathBuf};
+use std::sync::mpsc;
+use std::thread;
 
 const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
 const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
@@ -32,6 +36,9 @@ const GDN_CONV_BYTES: u64 = 2_211_840;
 const PLE_CONV_BYTES: u64 = 184_320;
 const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4;
 const MTP_CAPTURE_LOGITS_BYTES: u64 = 248_320 * 4;
+const PLE_ROW_BYTES: usize = 100;
+const PLE_READERS: usize = 16;
+const PLE_HOT_BYTES: usize = 1024 * 1024 * 1024;
 
 #[derive(Deserialize)]
 struct Manifest {
@@ -104,7 +111,9 @@ pub(super) struct ArtifactBindings {
 
 pub(super) struct QwenMap {
     path: PathBuf,
-    map: Mmap,
+    file: File,
+    bytes: u64,
+    map: Option<Mmap>,
 }
 
 #[derive(Clone)]
@@ -123,9 +132,227 @@ pub(super) struct QwenModel {
     memory: MemoryPlan,
     maps: Vec<QwenMap>,
     tensors: HashMap<String, QwenTensor>,
+    ple_reader: PleReader,
     identity: [u8; 32],
 }
 
+struct PleTask {
+    row: u64,
+    response: mpsc::Sender<(u64, Result<[u8; PLE_ROW_BYTES], String>)>,
+}
+
+struct PleReader {
+    workers: Vec<mpsc::Sender<Option<PleTask>>>,
+    threads: Vec<thread::JoinHandle<()>>,
+    hot: HashMap<u64, ([u8; PLE_ROW_BYTES], u64)>,
+    order: VecDeque<(u64, u64)>,
+    generation: u64,
+    capacity: usize,
+    hits: u64,
+    misses: u64,
+    evictions: u64,
+    read_bytes: u64,
+    read_ns: u64,
+}
+
+#[derive(Clone, Copy, Default)]
+pub(super) struct PleStats {
+    pub(super) entries: u64,
+    pub(super) cache_bytes: u64,
+    pub(super) hits: u64,
+    pub(super) misses: u64,
+    pub(super) evictions: u64,
+    pub(super) read_bytes: u64,
+    pub(super) read_ms: u64,
+}
+
+impl PleReader {
+    fn new(maps: &[QwenMap], tensors: &HashMap<String, QwenTensor>) -> Result<Self, String> {
+        let packed = tensors
+            .get("ngram.weight")
+            .ok_or_else(|| "Qwen tensor is missing: ngram.weight".to_owned())?;
+        let scales = tensors
+            .get("ngram.scales")
+            .ok_or_else(|| "Qwen tensor is missing: ngram.scales".to_owned())?;
+        let biases = tensors
+            .get("ngram.biases")
+            .ok_or_else(|| "Qwen tensor is missing: ngram.biases".to_owned())?;
+        if scales.map != packed.map || biases.map != packed.map {
+            return Err("Qwen PLE tensors do not share their sidecar file".into());
+        }
+        let ranges = [
+            (packed.range.clone(), 80_u64),
+            (scales.range.clone(), 10_u64),
+            (biases.range.clone(), 10_u64),
+        ];
+        let mut workers = Vec::with_capacity(PLE_READERS);
+        let mut threads = Vec::with_capacity(PLE_READERS);
+        for index in 0..PLE_READERS {
+            let file = maps[packed.map]
+                .file
+                .try_clone()
+                .map_err(|error| format!("cloning the Qwen PLE sidecar: {error}"))?;
+            let ranges = ranges.clone();
+            let (sender, receiver) = mpsc::channel::<Option<PleTask>>();
+            let handle = thread::Builder::new()
+                .name(format!("qwen-ple-{index}"))
+                .spawn(move || {
+                    while let Ok(Some(task)) = receiver.recv() {
+                        let result = read_ple_row(&file, &ranges, task.row);
+                        let _ = task.response.send((task.row, result));
+                    }
+                })
+                .map_err(|error| format!("starting a Qwen PLE reader: {error}"))?;
+            workers.push(sender);
+            threads.push(handle);
+        }
+        Ok(Self {
+            workers,
+            threads,
+            hot: HashMap::default(),
+            order: VecDeque::new(),
+            generation: 0,
+            capacity: PLE_HOT_BYTES / PLE_ROW_BYTES,
+            hits: 0,
+            misses: 0,
+            evictions: 0,
+            read_bytes: 0,
+            read_ns: 0,
+        })
+    }
+
+    fn rows(&mut self, rows: &[u64]) -> Result<Vec<[u8; PLE_ROW_BYTES]>, String> {
+        let mut unique = Vec::with_capacity(rows.len());
+        let mut seen = HashSet::with_capacity(rows.len());
+        for &row in rows {
+            if seen.insert(row) {
+                unique.push(row);
+            }
+        }
+        let missing = unique
+            .iter()
+            .copied()
+            .filter(|row| !self.hot.contains_key(row))
+            .collect::<Vec<_>>();
+        self.hits = self
+            .hits
+            .saturating_add((unique.len() - missing.len()) as u64);
+        self.misses = self.misses.saturating_add(missing.len() as u64);
+        if !missing.is_empty() {
+            let started = std::time::Instant::now();
+            let (response, received) = mpsc::channel();
+            for (index, row) in missing.iter().copied().enumerate() {
+                self.workers[index % self.workers.len()]
+                    .send(Some(PleTask {
+                        row,
+                        response: response.clone(),
+                    }))
+                    .map_err(|_| "Qwen PLE reader stopped unexpectedly".to_owned())?;
+            }
+            drop(response);
+            for _ in 0..missing.len() {
+                let (row, result) = received
+                    .recv()
+                    .map_err(|_| "Qwen PLE reader returned too few rows".to_owned())?;
+                self.insert(row, result?);
+            }
+            self.read_bytes = self
+                .read_bytes
+                .saturating_add((missing.len() * PLE_ROW_BYTES) as u64);
+            self.read_ns = self
+                .read_ns
+                .saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
+        }
+        rows.iter()
+            .map(|row| {
+                let value = self
+                    .hot
+                    .get(row)
+                    .map(|(value, _)| *value)
+                    .ok_or_else(|| format!("Qwen PLE row {row} was not cached"))?;
+                self.touch(*row);
+                Ok(value)
+            })
+            .collect()
+    }
+
+    fn insert(&mut self, row: u64, value: [u8; PLE_ROW_BYTES]) {
+        self.generation = self.generation.wrapping_add(1);
+        self.hot.insert(row, (value, self.generation));
+        self.order.push_back((row, self.generation));
+        while self.hot.len() > self.capacity {
+            let Some((old, generation)) = self.order.pop_front() else {
+                break;
+            };
+            if self
+                .hot
+                .get(&old)
+                .is_some_and(|entry| entry.1 == generation)
+            {
+                self.hot.remove(&old);
+                self.evictions = self.evictions.saturating_add(1);
+            }
+        }
+    }
+
+    fn touch(&mut self, row: u64) {
+        self.generation = self.generation.wrapping_add(1);
+        if let Some(entry) = self.hot.get_mut(&row) {
+            entry.1 = self.generation;
+            self.order.push_back((row, self.generation));
+        }
+    }
+
+    fn stats(&self) -> PleStats {
+        PleStats {
+            entries: self.hot.len() as u64,
+            cache_bytes: (self.hot.len() * PLE_ROW_BYTES) as u64,
+            hits: self.hits,
+            misses: self.misses,
+            evictions: self.evictions,
+            read_bytes: self.read_bytes,
+            read_ms: self.read_ns / 1_000_000,
+        }
+    }
+}
+
+impl Drop for PleReader {
+    fn drop(&mut self) {
+        for worker in &self.workers {
+            let _ = worker.send(None);
+        }
+        for thread in self.threads.drain(..) {
+            let _ = thread.join();
+        }
+    }
+}
+
+fn read_ple_row(
+    file: &File,
+    ranges: &[(std::ops::Range<u64>, u64); 3],
+    row: u64,
+) -> Result<[u8; PLE_ROW_BYTES], String> {
+    let mut output = [0_u8; PLE_ROW_BYTES];
+    let mut cursor = 0;
+    for (range, width) in ranges {
+        let offset = range
+            .start
+            .checked_add(
+                row.checked_mul(*width)
+                    .ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?,
+            )
+            .ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?;
+        if offset + width > range.end {
+            return Err(format!("Qwen PLE row {row} is truncated"));
+        }
+        let end = cursor + *width as usize;
+        file.read_exact_at(&mut output[cursor..end], offset)
+            .map_err(|error| format!("reading Qwen PLE row {row}: {error}"))?;
+        cursor = end;
+    }
+    Ok(output)
+}
+
 impl QwenModel {
     #[cfg(test)]
     pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
@@ -150,14 +377,28 @@ impl QwenModel {
         paths.sort();
         paths.dedup();
         let mut maps = Vec::with_capacity(paths.len());
-        let mut map_indices = HashMap::with_capacity(paths.len());
+        let mut map_indices = HashMap::with_capacity_and_hasher(paths.len(), Default::default());
         for path in paths {
             let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?;
             // SAFETY: verified managed artifacts remain read-only while the model owns each mapping.
             let map = unsafe { MmapOptions::new().map(&file) }
                 .map_err(|error| format!("cannot map {}: {error}", path.display()))?;
+            if path
+                .file_name()
+                .is_some_and(|name| name == "ngram-table.safetensors")
+            {
+                // SAFETY: the read-only mapping remains valid for this advisory call.
+                unsafe {
+                    libc::madvise(map.as_ptr().cast_mut().cast(), map.len(), libc::MADV_RANDOM);
+                }
+            }
             map_indices.insert(path.clone(), maps.len());
-            maps.push(QwenMap { path, map });
+            maps.push(QwenMap {
+                path,
+                file,
+                bytes: map.len() as u64,
+                map: Some(map),
+            });
         }
         let tensors = bindings
             .into_iter()
@@ -174,6 +415,7 @@ impl QwenModel {
                 (binding.name, tensor)
             })
             .collect::<HashMap<_, _>>();
+        let ple_reader = PleReader::new(&maps, &tensors)?;
         let mut hash = Sha256::new();
         hash.update(b"DS4Server Qwen3.8 checkpoint identity v1");
         hash.update(MANIFEST);
@@ -183,6 +425,7 @@ impl QwenModel {
             memory: loaded.memory,
             maps,
             tensors,
+            ple_reader,
             identity,
         })
     }
@@ -193,12 +436,25 @@ impl QwenModel {
             .ok_or_else(|| format!("Qwen tensor is missing: {name}"))
     }
 
+    pub(super) fn tensors(&self) -> impl Iterator<Item = &QwenTensor> {
+        self.tensors.values()
+    }
+
     pub(super) fn map(&self, index: usize) -> (&[u8], &Path) {
-        (&self.maps[index].map, &self.maps[index].path)
+        (
+            self.maps[index]
+                .map
+                .as_deref()
+                .expect("materialized Qwen weights must not use their released source mapping"),
+            &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 map = self.maps[tensor.map]
+            .map
+            .as_deref()
+            .ok_or_else(|| format!("{} source mapping was already released", tensor.name))?;
         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)
@@ -207,6 +463,48 @@ impl QwenModel {
             .ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name))
     }
 
+    pub(super) fn read_tensor_at(
+        &self,
+        tensor: &QwenTensor,
+        offset: u64,
+        output: &mut [u8],
+    ) -> Result<(), String> {
+        let bytes = tensor.range.end - tensor.range.start;
+        let end = offset
+            .checked_add(output.len() as u64)
+            .ok_or_else(|| format!("{} read offset overflows", tensor.name))?;
+        if end > bytes {
+            return Err(format!("{} read exceeds its tensor range", tensor.name));
+        }
+        self.maps[tensor.map]
+            .file
+            .read_exact_at(output, tensor.range.start + offset)
+            .map_err(|error| format!("reading {} at byte {offset}: {error}", tensor.name))
+    }
+
+    pub(super) fn release_materialized_mappings(&mut self) {
+        for (index, map) in self.maps.iter_mut().enumerate() {
+            if !self
+                .tensors
+                .values()
+                .any(|tensor| tensor.map == index && tensor.name.starts_with("ngram."))
+            {
+                map.map = None;
+            }
+        }
+    }
+
+    pub(super) fn read_ple_rows(
+        &mut self,
+        rows: &[u64],
+    ) -> Result<Vec<[u8; PLE_ROW_BYTES]>, String> {
+        self.ple_reader.rows(rows)
+    }
+
+    pub(super) fn ple_stats(&self) -> PleStats {
+        self.ple_reader.stats()
+    }
+
     pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
         self.identity
     }
@@ -214,7 +512,7 @@ impl QwenModel {
     pub(super) fn summary(&self) -> ModelSummary {
         ModelSummary {
             model: ModelChoice::Qwen38FlashNext,
-            mapped_bytes: self.maps.iter().map(|item| item.map.len() as u64).sum(),
+            mapped_bytes: self.maps.iter().map(|item| item.bytes).sum(),
             tensor_count: self.tensors.len(),
             vocabulary_size: self.tokenizer.vocab_size(),
             support_loaded: false,
@@ -284,12 +582,15 @@ impl QwenModel {
         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)];
+            let Some(map) = &item.map else {
+                continue;
+            };
+            let mut pages = vec![0_i8; 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(),
+                    map.as_ptr().cast_mut().cast(),
+                    map.len(),
                     pages.as_mut_ptr(),
                 )
             } != 0
@@ -301,7 +602,7 @@ impl QwenModel {
                 ));
             }
             let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
-                .min(item.map.len()) as u64;
+                .min(map.len()) as u64;
             if item
                 .path
                 .file_name()
@@ -846,8 +1147,8 @@ mod tests {
         assert_eq!(
             direct,
             [
-                248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 13_314, 741, 29,
-                271, 510, 26_003, 29, 271,
+                248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 248_068, 271,
+                248_069, 271,
             ]
         );
         assert_eq!(

src/engine/tokenizer.rs

@@ -62,7 +62,6 @@ struct QwenTokenizerFile {
 struct QwenAddedToken {
     id: usize,
     content: String,
-    special: bool,
 }
 
 #[derive(Deserialize)]
@@ -225,9 +224,7 @@ impl Tokenizer {
                 ));
             }
             tokens[token.id] = token.content.as_bytes().to_vec();
-            if token.special {
-                rendered_specials.push((token.content.into_bytes(), token.id as i32));
-            }
+            rendered_specials.push((token.content.into_bytes(), token.id as i32));
         }
         if tokens[..=maximum].iter().any(Vec::is_empty) {
             return Err("Qwen tokenizer token ids are not contiguous".into());