feat: add layer_devices field to Context for future multi-GPU pipeline parallelism

Adds the infrastructure (per-layer device map) but does not activate it yet —
naive layer splitting across GPUs is slower due to cross-device transfer overhead.
Proper multi-GPU requires pipeline parallelism in the forward pass where the
hidden state is explicitly moved between devices at layer boundaries.

Also includes parallel warmup + dequant optimizations from the flash-moe session.

Author: evilsocket

cake-core/src/cake/mod.rs

@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     fmt::{Debug, Display},
     path::PathBuf,
     sync::{Arc, Mutex},
@@ -58,6 +59,9 @@ pub struct Context {
     /// Tensor storage for expert offloading (pread from safetensors).
     /// None when expert_offload is disabled.
     pub tensor_storage: Option<Arc<dyn utils::tensor_storage::TensorStorageProvider>>,
+    /// Per-layer device overrides for multi-GPU support.
+    /// When set, layers are split across GPUs instead of using `self.device` for all.
+    pub layer_devices: Option<HashMap<String, Device>>,
 }
 
 /// Parse a dtype string ("f16", "bf16", "f32") into a candle DType.
@@ -497,6 +501,7 @@ impl Context {
             listener_override: Arc::new(Mutex::new(None)),
             backend,
             tensor_storage,
+            layer_devices: None,
         })
     }
 }

cake-core/src/cake/sharding/api/test_helpers.rs

@@ -282,6 +282,7 @@ fn dummy_context() -> Context {
         quant: std::sync::Arc::new(crate::utils::NoQuantization),
         listener_override: std::sync::Mutex::new(None).into(),
             tensor_storage: None,
+            layer_devices: None,
         backend: std::sync::Arc::new(crate::backends::CpuBackend::new()),
     }
 }

cake-core/src/cake/sharding/worker.rs

@@ -647,6 +647,7 @@ mod tests {
             quant: Arc::new(crate::utils::NoQuantization),
             listener_override: Arc::new(std::sync::Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
             backend: Arc::new(crate::backends::CpuBackend::new()),
         };
 
@@ -724,6 +725,7 @@ mod tests {
             quant: Arc::new(crate::utils::NoQuantization),
             listener_override: Arc::new(std::sync::Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
             backend: Arc::new(crate::backends::CpuBackend::new()),
         };
 

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

@@ -713,6 +713,7 @@ mod tests {
             quant: Arc::new(crate::utils::NoQuantization),
             listener_override: Arc::new(Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
             backend: Arc::new(crate::backends::CpuBackend::new()),
         }
     }

cake-core/src/models/qwen3_5_moe/block.rs

@@ -86,6 +86,7 @@ impl Forwarder for Qwen3_5MoeBlock {
         let moe = if let Some(storage) = &ctx.tensor_storage {
             let layer_prefix = format!("{name}.mlp");
             let mlp_vb = vb.pp("mlp");
+            let layer_device = ctx.device.clone();
             // Detect stacked switch_mlp format vs individual experts
             let switch_mlp_name = format!("{layer_prefix}.switch_mlp.gate_proj.weight");
             let has_switch_mlp = storage.has_tensor(&switch_mlp_name)
@@ -95,12 +96,12 @@ impl Forwarder for Qwen3_5MoeBlock {
                 // Stacked switch_mlp format — use disk provider with switch_mlp namespace
                 std::sync::Arc::new(crate::models::common::disk_expert_provider::DiskExpertProvider::new_stacked(
                     storage.clone(), format!("{layer_prefix}.switch_mlp"), cfg.num_experts,
-                    ctx.device.clone(), ctx.dtype,
+                    layer_device, ctx.dtype,
                 ))
             } else {
                 // Individual experts: stream from disk
                 std::sync::Arc::new(crate::models::common::disk_expert_provider::DiskExpertProvider::new(
-                    storage.clone(), layer_prefix, cfg.num_experts, ctx.device.clone(), ctx.dtype,
+                    storage.clone(), layer_prefix, cfg.num_experts, layer_device, ctx.dtype,
                     ctx.quant.gptq_group_size(),
                 ))
             };

cake-core/src/models/sd/util.rs

@@ -314,6 +314,7 @@ mod tests {
             quant: Arc::new(crate::utils::NoQuantization),
             listener_override: Arc::new(Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
             backend: Arc::new(crate::backends::CpuBackend::new()),
         }
     }

cake-core/src/models/vibevoice/vibevoice.rs

@@ -163,6 +163,7 @@ impl VibeVoiceTTS {
                 quant: std::sync::Arc::new(crate::utils::NoQuantization),
                 listener_override: std::sync::Arc::new(std::sync::Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
                 backend: crate::backends::create_backend(device),
             },
         })

cake-core/src/models/vibevoice/vibevoice_1_5b.rs

@@ -279,6 +279,7 @@ impl VibeVoice1_5B {
                 quant: std::sync::Arc::new(crate::utils::NoQuantization),
                 listener_override: std::sync::Arc::new(std::sync::Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
                 backend: crate::backends::create_backend(device),
             },
         })

cake-core/tests/unit_tests/test_blocks.rs

@@ -36,6 +36,7 @@ fn make_context(cfg: Config, vb: VarBuilder<'static>) -> Context {
         quant: Arc::new(cake_core::utils::NoQuantization),
         listener_override: Arc::new(Mutex::new(None)),
             tensor_storage: None,
+            layer_devices: None,
         backend: Arc::new(cake_core::backends::CpuBackend::new()),
     }
 }