Finish DS4 execution parity

src/engine/metal/gpu.rs

@@ -18,6 +18,23 @@ pub(super) struct StreamExpertTable {
     pub(super) down_expert_bytes: u64,
 }
 
+#[derive(Clone, Copy, Default)]
+#[repr(C)]
+pub(super) struct StreamExpertCacheStats {
+    pub(super) configured_count: u32,
+    pub(super) current_count: u32,
+    pub(super) hits: u64,
+    pub(super) misses: u64,
+    pub(super) evictions: u64,
+    pub(super) wraps: u64,
+    pub(super) buffer_allocs: u64,
+    pub(super) buffer_reuses: u64,
+    pub(super) evict_advise_bytes: u64,
+    pub(super) willneed_advise_bytes: u64,
+    pub(super) pread_bytes: u64,
+    pub(super) pread_ms: f64,
+}
+
 unsafe extern "C" {
     pub(super) fn ds4_gpu_init() -> i32;
     pub(super) fn ds4_gpu_cleanup();
@@ -43,10 +60,7 @@ unsafe extern "C" {
     pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
     pub(super) fn ds4_gpu_set_streaming_expert_cache_expert_bytes(bytes: u64);
     pub(super) fn ds4_gpu_recommended_working_set_size() -> u64;
-    pub(super) fn ds4_gpu_stream_expert_cache_budget_for_expert_size(
-        gate_expert_bytes: u64,
-        down_expert_bytes: u64,
-    ) -> u32;
+    pub(super) fn ds4_gpu_stream_expert_cache_get_stats(stats: *mut StreamExpertCacheStats);
     pub(super) fn ds4_gpu_stream_expert_cache_seed_experts(
         table: *const StreamExpertTable,
         expert_ids: *const i32,
@@ -266,6 +280,17 @@ unsafe extern "C" {
         token: u32,
         embd: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_embed_tokens_quant_tensor(
+        out: *mut GpuTensor,
+        tokens: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        vocab: u32,
+        rows: u32,
+        embd: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_rope_tail_tensor(
         x: *mut GpuTensor,
         tokens: u32,
@@ -301,6 +326,30 @@ unsafe extern "C" {
         cache_f16: bool,
         eps: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_kv_lora_rms_norm_tensor(
+        out: *mut GpuTensor,
+        kv_raw: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        rows: u32,
+        kv_raw_dim: u32,
+        kv_lora: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_store_compact_kv_tensor(
+        kv_cache: *mut GpuTensor,
+        rope_cache: *mut GpuTensor,
+        kv_norm: *const GpuTensor,
+        kv_raw: *const GpuTensor,
+        pos: u32,
+        rows: u32,
+        cache_cap: u32,
+        kv_raw_dim: u32,
+        kv_lora: u32,
+        rot: u32,
+        cache_f16: bool,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_store_indexer_k_tensor(
         cache: *mut GpuTensor,
         raw: *const GpuTensor,
@@ -327,6 +376,13 @@ unsafe extern "C" {
         selected: *mut GpuTensor,
         count: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_fill_selected_range_batch_tensor(
+        selected: *mut GpuTensor,
+        rows: u32,
+        pos: u32,
+        count: u32,
+        pad_row: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_indexer_rope_tail_tensor(
         x: *mut GpuTensor,
         tokens: u32,
@@ -353,6 +409,19 @@ unsafe extern "C" {
         scale: f32,
         cache_f16: bool,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_indexer_scores_batch_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        cache: *const GpuTensor,
+        visible: u32,
+        rows: u32,
+        pos: u32,
+        heads: u32,
+        head_dim: u32,
+        scale: f32,
+        cache_f16: bool,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
         out: *mut GpuTensor,
         q: *const GpuTensor,
@@ -365,6 +434,31 @@ unsafe extern "C" {
         q_nope: u32,
         q_dim: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
+        out: *mut GpuTensor,
+        q: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        rows: u32,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+        q_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_value_project_typed_batch_heads_tensor(
+        heads: *mut GpuTensor,
+        lora: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        rows: u32,
+        n_head: u32,
+        kv_lora: u32,
+        value_dim: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_attention_indexed_decode_typed_tensor(
         heads_out: *mut GpuTensor,
         q: *const GpuTensor,
@@ -392,6 +486,95 @@ unsafe extern "C" {
         beta_fast: f32,
         beta_slow: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
+        out: *mut GpuTensor,
+        q: *const GpuTensor,
+        qk_low: *const GpuTensor,
+        kv_cache: *const GpuTensor,
+        rope_cache: *const GpuTensor,
+        rows: u32,
+        pos: u32,
+        selected: u32,
+        cache_cap: u32,
+        cache_f16: bool,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+        rot: u32,
+        original: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
+        out: *mut GpuTensor,
+        q: *const GpuTensor,
+        qk_low: *const GpuTensor,
+        kv_cache: *const GpuTensor,
+        rope_cache: *const GpuTensor,
+        selected: *const GpuTensor,
+        rows: u32,
+        selected_count: u32,
+        cache_cap: u32,
+        cache_f16: bool,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+        rot: u32,
+        original: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_router_select_batch_tensor(
+        selected: *mut GpuTensor,
+        weights: *mut GpuTensor,
+        probs: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        bias: u64,
+        logits: *const GpuTensor,
+        experts: u32,
+        experts_used: u32,
+        scale: f32,
+        rows: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_routed_moe_batch_tensor(
+        out: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate_weight: u64,
+        up_weight: u64,
+        down_weight: u64,
+        gate_type: u32,
+        up_type: u32,
+        down_type: u32,
+        gate_expert_bytes: u64,
+        gate_row_bytes: u64,
+        up_expert_bytes: u64,
+        up_row_bytes: u64,
+        down_expert_bytes: u64,
+        down_row_bytes: u64,
+        input: u32,
+        middle: u32,
+        output: u32,
+        selected: *const GpuTensor,
+        weights: *const GpuTensor,
+        total_experts: u32,
+        used_experts: u32,
+        layer: u32,
+        x: *const GpuTensor,
+        rows: u32,
+        mid_token_stride: u32,
+        force_resident: bool,
+    ) -> i32;
     pub(super) fn ds4_gpu_glm_router_select_tensor(
         selected: *mut GpuTensor,
         weights: *mut GpuTensor,
@@ -1276,6 +1459,20 @@ impl Buffer {
         )
     }
 
+    pub(super) fn read_i32(&self, values: &mut [i32]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_read(
+                    self.raw(),
+                    0,
+                    values.as_mut_ptr().cast(),
+                    std::mem::size_of_val(values) as u64,
+                )
+            },
+            "reading Metal integers",
+        )
+    }
+
     pub(super) fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
         check(
             unsafe {

src/engine/metal/profile.rs

@@ -0,0 +1,307 @@
+use super::gpu::Buffer;
+use crate::model::ModelChoice;
+use serde_json::{Value, json};
+use std::cmp::Ordering;
+use std::fs::{self, File};
+use std::path::PathBuf;
+
+const CACHE_CAPS: [usize; 10] = [1, 2, 4, 8, 16, 32, 64, 128, 256, 384];
+
+#[derive(Default)]
+struct LayerProfile {
+    records: u64,
+    counts: Vec<u64>,
+    weights: Vec<f64>,
+    caches: Vec<Vec<i32>>,
+    cache_hits: Vec<u64>,
+    cache_weight_hits: Vec<f64>,
+    previous: Option<(u32, Vec<i32>)>,
+    adjacent_pairs: u64,
+    adjacent_overlap: f64,
+    adjacent_jaccard: f64,
+    hash_router: bool,
+}
+
+pub(super) struct ExpertProfile {
+    path: PathBuf,
+    model: ModelChoice,
+    experts: usize,
+    used: usize,
+    caps: Vec<usize>,
+    layers: Vec<LayerProfile>,
+}
+
+impl ExpertProfile {
+    pub(super) fn new(
+        path: Option<&str>,
+        model: ModelChoice,
+        layers: u32,
+        experts: u64,
+        used: u64,
+    ) -> Result<Option<Self>, String> {
+        let Some(path) = path else { return Ok(None) };
+        if path.trim().is_empty() {
+            return Err("Expert profile path cannot be empty".into());
+        }
+        let experts = usize::try_from(experts).map_err(|_| "expert count is too large")?;
+        let used = usize::try_from(used).map_err(|_| "selected expert count is too large")?;
+        let caps = CACHE_CAPS
+            .into_iter()
+            .filter(|cap| *cap <= experts)
+            .collect::<Vec<_>>();
+        let layers = (0..layers)
+            .map(|_| LayerProfile {
+                counts: vec![0; experts],
+                weights: vec![0.0; experts],
+                caches: vec![Vec::new(); caps.len()],
+                cache_hits: vec![0; caps.len()],
+                cache_weight_hits: vec![0.0; caps.len()],
+                ..LayerProfile::default()
+            })
+            .collect();
+        Ok(Some(Self {
+            path: path.into(),
+            model,
+            experts,
+            used,
+            caps,
+            layers,
+        }))
+    }
+
+    pub(super) fn record(
+        &mut self,
+        layer: usize,
+        pos: u32,
+        selected: &Buffer,
+        weights: &Buffer,
+        rows: u32,
+        hash_router: bool,
+    ) -> Result<(), String> {
+        let count = self
+            .used
+            .checked_mul(rows as usize)
+            .ok_or("expert profile row count overflow")?;
+        let mut ids = vec![0; count];
+        let mut route_weights = vec![0.0; count];
+        super::call(
+            unsafe { super::gpu::ds4_gpu_end_commands() },
+            "ending Metal commands for expert profiling",
+        )?;
+        let read = selected
+            .read_i32(&mut ids)
+            .and_then(|()| weights.read_f32(&mut route_weights));
+        let resumed = super::call(
+            unsafe { super::gpu::ds4_gpu_begin_commands() },
+            "resuming Metal commands after expert profiling",
+        );
+        read.and(resumed)?;
+        for (row, (ids, weights)) in ids
+            .chunks_exact(self.used)
+            .zip(route_weights.chunks_exact(self.used))
+            .enumerate()
+        {
+            self.record_row(layer, pos + row as u32, ids, weights, hash_router)?;
+        }
+        Ok(())
+    }
+
+    fn record_row(
+        &mut self,
+        layer: usize,
+        pos: u32,
+        ids: &[i32],
+        weights: &[f32],
+        hash_router: bool,
+    ) -> Result<(), String> {
+        let profile = self
+            .layers
+            .get_mut(layer)
+            .ok_or("expert profile layer is outside the model")?;
+        profile.records += 1;
+        profile.hash_router |= hash_router;
+        if let Some((previous_pos, previous)) = &profile.previous
+            && previous_pos.checked_add(1) == Some(pos)
+        {
+            let intersection = previous.iter().filter(|id| ids.contains(id)).count();
+            profile.adjacent_pairs += 1;
+            profile.adjacent_overlap += intersection as f64 / self.used as f64;
+            profile.adjacent_jaccard += intersection as f64 / (2 * self.used - intersection) as f64;
+        }
+        profile.previous = Some((pos, ids.to_vec()));
+        for (&id, &weight) in ids.iter().zip(weights) {
+            let expert = usize::try_from(id)
+                .ok()
+                .filter(|expert| *expert < self.experts)
+                .ok_or_else(|| format!("router selected invalid expert {id} at layer {layer}"))?;
+            profile.counts[expert] += 1;
+            profile.weights[expert] += f64::from(weight);
+            for (index, (&cap, cache)) in self.caps.iter().zip(&mut profile.caches).enumerate() {
+                if let Some(found) = cache.iter().position(|cached| *cached == id) {
+                    profile.cache_hits[index] += 1;
+                    profile.cache_weight_hits[index] += f64::from(weight);
+                    cache.remove(found);
+                } else if cache.len() == cap {
+                    cache.pop();
+                }
+                cache.insert(0, id);
+            }
+        }
+        Ok(())
+    }
+
+    pub(super) fn write(&self) -> Result<(), String> {
+        let selections: u64 = self
+            .layers
+            .iter()
+            .map(|layer| layer.counts.iter().sum::<u64>())
+            .sum();
+        let weight_total = self
+            .layers
+            .iter()
+            .flat_map(|layer| &layer.weights)
+            .sum::<f64>();
+        let cache_summary = self
+            .caps
+            .iter()
+            .enumerate()
+            .map(|(index, cap)| {
+                let hits = self
+                    .layers
+                    .iter()
+                    .map(|layer| layer.cache_hits[index])
+                    .sum::<u64>();
+                let weighted = self
+                    .layers
+                    .iter()
+                    .map(|layer| layer.cache_weight_hits[index])
+                    .sum::<f64>();
+                json!({
+                    "n": cap,
+                    "hits": hits,
+                    "selections": selections,
+                    "hit_rate": fraction(hits as f64, selections as f64),
+                    "weighted_hit_rate": fraction(weighted, weight_total),
+                })
+            })
+            .collect::<Vec<_>>();
+        let layers = self
+            .layers
+            .iter()
+            .enumerate()
+            .map(|(index, layer)| self.layer_json(index, layer))
+            .collect::<Vec<_>>();
+        let value = json!({
+            "source": "ds4 Metal expert locality profile",
+            "model": self.model.id(),
+            "layers": self.layers.len(),
+            "experts": self.experts,
+            "expert_used": self.used,
+            "layer_records": self.layers.iter().map(|layer| layer.records).sum::<u64>(),
+            "selections": selections,
+            "cache_ns": self.caps,
+            "cache_summary": cache_summary,
+            "layers_detail": layers,
+        });
+        if let Some(parent) = self
+            .path
+            .parent()
+            .filter(|parent| !parent.as_os_str().is_empty())
+        {
+            fs::create_dir_all(parent).map_err(|error| error.to_string())?;
+        }
+        let temporary = self.path.with_extension("tmp");
+        let file = File::create(&temporary).map_err(|error| error.to_string())?;
+        serde_json::to_writer_pretty(&file, &value).map_err(|error| error.to_string())?;
+        file.sync_all().map_err(|error| error.to_string())?;
+        fs::rename(temporary, &self.path).map_err(|error| error.to_string())
+    }
+
+    fn layer_json(&self, index: usize, layer: &LayerProfile) -> Value {
+        let selections = layer.counts.iter().sum::<u64>();
+        let total_weight = layer.weights.iter().sum::<f64>();
+        let mut experts = (0..self.experts)
+            .filter(|expert| layer.counts[*expert] != 0)
+            .collect::<Vec<_>>();
+        experts.sort_by(|a, b| {
+            layer.counts[*b]
+                .cmp(&layer.counts[*a])
+                .then_with(|| {
+                    layer.weights[*b]
+                        .partial_cmp(&layer.weights[*a])
+                        .unwrap_or(Ordering::Equal)
+                })
+                .then_with(|| a.cmp(b))
+        });
+        let top = experts
+            .into_iter()
+            .take(16)
+            .map(|expert| {
+                json!({
+                    "id": expert,
+                    "count": layer.counts[expert],
+                    "pct": 100.0 * fraction(layer.counts[expert] as f64, selections as f64),
+                    "weight": layer.weights[expert],
+                    "weight_pct": 100.0 * fraction(layer.weights[expert], total_weight),
+                })
+            })
+            .collect::<Vec<_>>();
+        let cache = self
+            .caps
+            .iter()
+            .enumerate()
+            .map(|(cap_index, cap)| {
+                json!({
+                    "n": cap,
+                    "hits": layer.cache_hits[cap_index],
+                    "hit_rate": fraction(layer.cache_hits[cap_index] as f64, selections as f64),
+                    "weighted_hit_rate": fraction(layer.cache_weight_hits[cap_index], total_weight),
+                })
+            })
+            .collect::<Vec<_>>();
+        json!({
+            "layer": index,
+            "hash_router": layer.hash_router,
+            "records": layer.records,
+            "selections": selections,
+            "unique_experts": layer.counts.iter().filter(|count| **count != 0).count(),
+            "avg_adjacent_overlap": fraction(layer.adjacent_overlap, layer.adjacent_pairs as f64),
+            "avg_adjacent_jaccard": fraction(layer.adjacent_jaccard, layer.adjacent_pairs as f64),
+            "top_experts": top,
+            "cache": cache,
+        })
+    }
+}
+
+fn fraction(value: f64, total: f64) -> f64 {
+    if total > 0.0 { value / total } else { 0.0 }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn profile_reports_locality_and_lru_hits() {
+        let path = std::env::temp_dir().join(format!(
+            "ds4-profile-{}-{}.json",
+            std::process::id(),
+            std::thread::current().name().unwrap_or("test")
+        ));
+        let mut profile = ExpertProfile::new(path.to_str(), ModelChoice::DeepSeekV4Flash, 1, 8, 2)
+            .unwrap()
+            .unwrap();
+        profile
+            .record_row(0, 10, &[1, 2], &[0.6, 0.4], false)
+            .unwrap();
+        profile
+            .record_row(0, 11, &[1, 3], &[0.7, 0.3], false)
+            .unwrap();
+        profile.write().unwrap();
+        let value: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap();
+        assert_eq!(value["selections"], 4);
+        assert_eq!(value["layers_detail"][0]["avg_adjacent_overlap"], 0.5);
+        assert_eq!(value["cache_summary"][1]["hits"], 1);
+        let _ = fs::remove_file(path);
+    }
+}