flash-moe: expert cache + batched transfers — 16.6× speedup (0.16 → 2.66 tok/s)

- LRU cache for dequantized expert weights (avoids re-dequantization)
- Pre-warm cache at load time (moves dequant cost to model loading)
- Batch gate+up CPU→GPU transfer (1 PCIe transfer instead of 2)
- RwLock for concurrent cache reads
- Memory: ~60 GiB RAM for all cached experts (trades RAM for speed)

Author: evilsocket

cake-core/src/models/common/disk_expert_provider.rs

@@ -8,14 +8,22 @@
 //! Memory usage: O(num_experts_per_tok × expert_size) for the buffer pool,
 //! plus whatever the OS decides to keep in the page cache.
 
-use std::sync::Arc;
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
 
 use candle_core::{DType, Device, Result, Tensor};
 
 use crate::utils::tensor_storage::TensorStorageProvider;
 
 use super::expert_provider::{ExpertProvider, ExpertWeights};
 
+/// Simple LRU cache for dequantized expert weights.
+/// Avoids re-dequantizing the same popular experts across tokens.
+struct ExpertCache {
+    entries: RwLock<HashMap<usize, ExpertWeights>>,
+    capacity: usize,
+}
+
 /// Pre-computed tensor names for a single expert (avoids format! on hot path).
 struct ExpertNames {
     gate_proj: String,
@@ -86,6 +94,8 @@ pub struct DiskExpertProvider {
     stacked_meta: Option<StackedMeta>,
     /// Pre-computed stacked quantization tensor names.
     stacked_quant_names: Option<StackedQuantNames>,
+    /// Cache of dequantized expert weights (avoids repeated dequantization for popular experts).
+    cache: Option<ExpertCache>,
 }
 
 impl std::fmt::Debug for DiskExpertProvider {
@@ -170,6 +180,7 @@ impl DiskExpertProvider {
             gptq_group_size,
             stacked_meta: None,
             stacked_quant_names: None,
+            cache: None,
         }
     }
 
@@ -272,7 +283,7 @@ impl DiskExpertProvider {
         let storage_dtype = stacked_meta.as_ref().map(|m| m.storage_dtype);
         let use_f32_zerocopy = !is_affine && dtype == DType::F32
             && storage_dtype.is_some_and(|sd| sd == DType::F32);
-        Self {
+        let provider = Self {
             storage,
             layer_prefix,
             expert_names,
@@ -284,7 +295,33 @@ impl DiskExpertProvider {
             gptq_group_size: None,
             stacked_meta,
             stacked_quant_names,
+            // Enable cache for quantized experts — dequantization is expensive
+            cache: if is_affine {
+                Some(ExpertCache {
+                    entries: RwLock::new(HashMap::with_capacity(num_experts)),
+                    capacity: num_experts,
+                })
+            } else {
+                None
+            },
+        };
+
+        // Pre-warm: dequantize all experts at construction (moves cost from first token to loading)
+        if provider.cache.is_some() {
+            log::info!("pre-warming expert cache for {} experts...", num_experts);
+            for i in 0..num_experts {
+                if let Ok(ew) = provider.get_expert_uncached(i) {
+                    if let Some(ref cache) = provider.cache {
+                        if let Ok(mut entries) = cache.entries.write() {
+                            entries.insert(i, ew);
+                        }
+                    }
+                }
+            }
+            log::info!("expert cache warmed ({} entries)", num_experts);
         }
+
+        provider
     }
 
     /// Reinterpret a `Vec<u8>` as `Vec<f32>` without copying.
@@ -410,6 +447,87 @@ impl ExpertProvider for DiskExpertProvider {
             )));
         }
 
+        // Check cache first — stores CPU-side dequantized tensors to avoid repeated dequant
+        if let Some(ref cache) = self.cache {
+            if let Ok(entries) = cache.entries.read() {
+                if let Some(ew) = entries.get(&idx) {
+                    // Cache hit: transfer CPU tensors to target device
+                    return if self.needs_device_transfer {
+                        Ok(ExpertWeights {
+                            gate_proj: ew.gate_proj.to_device(&self.device)?,
+                            up_proj: ew.up_proj.to_device(&self.device)?,
+                            down_proj: ew.down_proj.to_device(&self.device)?,
+                        })
+                    } else {
+                        Ok(ew.clone())
+                    };
+                }
+            }
+        }
+
+        // get_expert_uncached returns CPU tensors when cache is enabled (stacked affine path)
+        let cpu_result = self.get_expert_uncached(idx)?;
+
+        // Store CPU tensors in cache
+        if let Some(ref cache) = self.cache {
+            if let Ok(mut entries) = cache.entries.write() {
+                if entries.len() < cache.capacity {
+                    entries.insert(idx, cpu_result.clone());
+                }
+            }
+        }
+
+        // Transfer to target device — batch gate+up into one PCIe transfer
+        if self.needs_device_transfer {
+            let gate_up = Tensor::cat(&[&cpu_result.gate_proj, &cpu_result.up_proj], 0)?;
+            let gate_up_gpu = gate_up.to_device(&self.device)?;
+            let g_rows = cpu_result.gate_proj.dim(0)?;
+            Ok(ExpertWeights {
+                gate_proj: gate_up_gpu.narrow(0, 0, g_rows)?,
+                up_proj: gate_up_gpu.narrow(0, g_rows, g_rows)?,
+                down_proj: cpu_result.down_proj.to_device(&self.device)?,
+            })
+        } else {
+            Ok(cpu_result)
+        }
+    }
+
+    fn num_experts(&self) -> usize {
+        self.num_experts
+    }
+
+    #[cfg(unix)]
+    fn prefetch_experts(&self, indices: &[usize]) {
+        for &idx in indices {
+            if idx >= self.num_experts {
+                continue;
+            }
+            let names = &self.expert_names[idx];
+            for name in [&names.gate_proj, &names.up_proj, &names.down_proj] {
+                if let Some((bytes, _, _)) = self.storage.tensor_bytes(name) {
+                    unsafe {
+                        libc::posix_madvise(
+                            bytes.as_ptr() as *mut _,
+                            bytes.len(),
+                            libc::POSIX_MADV_WILLNEED,
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl DiskExpertProvider {
+    /// Internal: load expert weights without cache lookup.
+    fn get_expert_uncached(&self, idx: usize) -> Result<ExpertWeights> {
+        if idx >= self.num_experts {
+            return Err(candle_core::Error::Msg(format!(
+                "expert index {idx} out of range (num_experts={})",
+                self.num_experts
+            )));
+        }
+
         let names = &self.expert_names[idx];
 
         // GPTQ path: read and dequantize each weight individually
@@ -443,8 +561,8 @@ impl ExpertProvider for DiskExpertProvider {
                     let scales = Tensor::from_raw_buffer(s_bytes, DType::BF16, &proj.scales_shape, &Device::Cpu)?;
                     let biases = Tensor::from_raw_buffer(b_bytes, DType::BF16, &proj.scales_shape, &Device::Cpu)?;
                     let weight = crate::utils::gptq::dequantize_packed_4bit(&packed, &scales, &biases, sm.group_size)?;
-                    let weight = weight.to_dtype(self.dtype)?;
-                    if self.needs_device_transfer { weight.to_device(target_device) } else { Ok(weight) }
+                    // Return CPU tensor — device transfer happens in get_expert after caching
+                    weight.to_dtype(self.dtype)
                 };
                 return Ok(ExpertWeights {
                     gate_proj: read_dequant(&names.gate_proj, &sm.gate, &qn.gate_scales, &qn.gate_biases)?,
@@ -563,31 +681,6 @@ impl ExpertProvider for DiskExpertProvider {
         })
     }
 
-    fn num_experts(&self) -> usize {
-        self.num_experts
-    }
-
-    #[cfg(unix)]
-    fn prefetch_experts(&self, indices: &[usize]) {
-        for &idx in indices {
-            if idx >= self.num_experts {
-                continue;
-            }
-            let names = &self.expert_names[idx];
-            // Issue madvise(WILLNEED) for each projection's mmap region
-            for name in [&names.gate_proj, &names.up_proj, &names.down_proj] {
-                if let Some((bytes, _, _)) = self.storage.tensor_bytes(name) {
-                    unsafe {
-                        libc::posix_madvise(
-                            bytes.as_ptr() as *mut _,
-                            bytes.len(),
-                            libc::POSIX_MADV_WILLNEED,
-                        );
-                    }
-                }
-            }
-        }
-    }
 }
 
 #[cfg(test)]