Plan split topology with exact layer weights

crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs

@@ -1207,6 +1207,7 @@ mod tests {
                 .to_string(),
             source_model_bytes: 1234,
             source_files: Vec::new(),
+            layer_weight_bytes: Vec::new(),
             layer_count,
             activation_width: 4096,
             tensor_count: 100,

crates/mesh-llm-host-runtime/src/inference/skippy/package.rs

@@ -17,6 +17,7 @@ pub struct SkippyPackageIdentity {
     pub source_model_sha256: String,
     pub source_model_bytes: u64,
     pub source_files: Vec<SkippyPackageSourceFile>,
+    pub layer_weight_bytes: Vec<u64>,
     pub layer_count: u32,
     pub activation_width: u32,
     pub tensor_count: u64,
@@ -108,6 +109,7 @@ pub fn synthetic_direct_gguf_package(
         source_model_sha256,
         source_model_bytes,
         source_files,
+        layer_weight_bytes: Vec::new(),
         layer_count: compact.layer_count,
         activation_width: compact.embedding_size,
         tensor_count,
@@ -315,6 +317,7 @@ pub fn identity_from_layer_package(package_ref: &str) -> Result<SkippyPackageIde
     let source_model_bytes = info
         .source_model_bytes
         .unwrap_or_else(|| info.layers.iter().map(|l| l.artifact_bytes).sum::<u64>());
+    let layer_weight_bytes = layer_weight_bytes_from_info(&info);
 
     // For local paths inside an HF cache, convert to an exact hf:// ref so all
     // nodes resolve the same snapshot independently. HF cache dirs look like:
@@ -328,13 +331,45 @@ pub fn identity_from_layer_package(package_ref: &str) -> Result<SkippyPackageIde
         source_model_sha256: info.source_model_sha256,
         source_model_bytes,
         source_files: Vec::new(),
+        layer_weight_bytes,
         layer_count: info.layer_count,
         activation_width,
         tensor_count: info.layers.iter().map(|l| l.tensor_count as u64).sum(),
         generation: info.generation,
     })
 }
 
+fn layer_weight_bytes_from_info(info: &skippy_runtime::package::LayerPackageInfo) -> Vec<u64> {
+    let mut layers = info.layers.clone();
+    layers.sort_by_key(|layer| layer.layer_index);
+    if layers.len() != info.layer_count as usize
+        || layers
+            .iter()
+            .enumerate()
+            .any(|(index, layer)| layer.layer_index as usize != index)
+    {
+        return Vec::new();
+    }
+    let mut weights = layers
+        .into_iter()
+        .map(|layer| layer.tensor_bytes.max(layer.artifact_bytes))
+        .collect::<Vec<_>>();
+    let accounted = weights.iter().copied().sum::<u64>();
+    let unaccounted = info
+        .source_model_bytes
+        .unwrap_or_default()
+        .saturating_sub(accounted);
+    if let Some((first, rest)) = weights.split_first_mut() {
+        *first = first.saturating_add(unaccounted.div_ceil(2));
+        if let Some(last) = rest.last_mut() {
+            *last = last.saturating_add(unaccounted / 2);
+        } else {
+            *first = first.saturating_add(unaccounted / 2);
+        }
+    }
+    weights
+}
+
 /// Detect if a local path is inside an HF cache directory and convert to `hf://` ref.
 ///
 /// HF cache paths look like:
@@ -521,4 +556,70 @@ mod tests {
         assert!(error.contains("missing activation_width"));
         assert!(error.contains("rebuild the package manifest"));
     }
+
+    #[test]
+    fn package_layer_weights_include_shared_model_bytes_at_endpoints() {
+        let info = skippy_runtime::package::LayerPackageInfo {
+            package_dir: PathBuf::from("/models/package"),
+            manifest_sha256: "manifest".to_string(),
+            model_id: "org/model".to_string(),
+            source_model_path: "model.gguf".to_string(),
+            source_model_sha256: "source".to_string(),
+            source_model_bytes: Some(120),
+            layer_count: 2,
+            activation_width: Some(1024),
+            generation: None,
+            projectors: Vec::new(),
+            layers: vec![
+                skippy_runtime::package::LayerPackageLayerInfo {
+                    layer_index: 0,
+                    tensor_count: 1,
+                    tensor_bytes: 30,
+                    artifact_bytes: 30,
+                },
+                skippy_runtime::package::LayerPackageLayerInfo {
+                    layer_index: 1,
+                    tensor_count: 1,
+                    tensor_bytes: 40,
+                    artifact_bytes: 40,
+                },
+            ],
+        };
+
+        assert_eq!(layer_weight_bytes_from_info(&info), vec![55, 65]);
+    }
+
+    #[test]
+    fn package_layer_weights_require_contiguous_indices() {
+        let mut info = skippy_runtime::package::LayerPackageInfo {
+            package_dir: PathBuf::from("/models/package"),
+            manifest_sha256: "manifest".to_string(),
+            model_id: "org/model".to_string(),
+            source_model_path: "model.gguf".to_string(),
+            source_model_sha256: "source".to_string(),
+            source_model_bytes: Some(70),
+            layer_count: 2,
+            activation_width: Some(1024),
+            generation: None,
+            projectors: Vec::new(),
+            layers: vec![
+                skippy_runtime::package::LayerPackageLayerInfo {
+                    layer_index: 0,
+                    tensor_count: 1,
+                    tensor_bytes: 30,
+                    artifact_bytes: 30,
+                },
+                skippy_runtime::package::LayerPackageLayerInfo {
+                    layer_index: 2,
+                    tensor_count: 1,
+                    tensor_bytes: 40,
+                    artifact_bytes: 40,
+                },
+            ],
+        };
+
+        assert!(layer_weight_bytes_from_info(&info).is_empty());
+        info.layers[1].layer_index = 1;
+        assert_eq!(layer_weight_bytes_from_info(&info), vec![30, 40]);
+    }
 }

crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs

@@ -73,6 +73,7 @@ pub(super) fn fake_package_identity(layer_count: u32) -> SkippyPackageIdentity {
             .to_string(),
         source_model_bytes: 1234,
         source_files: Vec::new(),
+        layer_weight_bytes: Vec::new(),
         layer_count,
         activation_width: 4096,
         tensor_count: 100,

crates/mesh-llm-host-runtime/src/runtime/local.rs

@@ -3844,6 +3844,7 @@ mod tests {
             source_model_sha256: "source".to_string(),
             source_model_bytes: u64::from(layer_count) * 1_000_000,
             source_files: Vec::new(),
+            layer_weight_bytes: Vec::new(),
             layer_count,
             activation_width: 2048,
             tensor_count: 100,

crates/mesh-llm-host-runtime/src/runtime/split_planning.rs

@@ -18,6 +18,7 @@ pub(super) struct SplitTopologyPlanInput {
     pub(super) native_context_length: u32,
     pub(super) layer_count: u32,
     pub(super) model_weight_bytes: u64,
+    pub(super) layer_weight_bytes: Vec<u64>,
     pub(super) kv_bytes_per_token: u64,
     pub(super) context_length_override: Option<u32>,
     pub(super) parallel_lanes_override: Option<usize>,
@@ -74,6 +75,7 @@ pub(super) fn plan_split_topology(input: SplitTopologyPlanInput) -> Result<Split
         native_context_length: input.native_context_length,
         layer_count: input.layer_count,
         model_weight_bytes: input.model_weight_bytes,
+        layer_weight_bytes: input.layer_weight_bytes,
         kv_bytes_per_token: input.kv_bytes_per_token,
         minimum_nodes: input.minimum_nodes,
         nodes: input
@@ -222,6 +224,7 @@ fn runtime_slice_plan_input(
         native_context_length: resources.native_context_length,
         layer_count: package.layer_count,
         model_weight_bytes: package.source_model_bytes,
+        layer_weight_bytes: package_layer_weight_bytes(package),
         kv_bytes_per_token: resources.kv_bytes_per_token,
         context_length_override: resources.ctx_size_override,
         parallel_lanes_override: resources.parallel_override,
@@ -240,6 +243,13 @@ fn runtime_slice_plan_input(
     }
 }
 
+fn package_layer_weight_bytes(package: &skippy::SkippyPackageIdentity) -> Vec<u64> {
+    if package.layer_weight_bytes.len() == package.layer_count as usize {
+        return package.layer_weight_bytes.clone();
+    }
+    Vec::new()
+}
+
 fn map_runtime_slice_stages(
     stages: Vec<TopologyStagePlan>,
     participant_by_id: &HashMap<String, SplitParticipant>,
@@ -548,6 +558,7 @@ mod tests {
             source_model_sha256: "source".to_string(),
             source_model_bytes,
             source_files: Vec::new(),
+            layer_weight_bytes: Vec::new(),
             layer_count,
             activation_width: 896,
             tensor_count: 100,
@@ -635,6 +646,39 @@ mod tests {
         assert_eq!(plan.stages.last().unwrap().layer_end, 30);
     }
 
+    #[test]
+    fn resource_planner_uses_exact_package_layer_weights() {
+        const GIB: u64 = 1024 * 1024 * 1024;
+        let participants = vec![participant(1, 12 * GIB), participant(2, 9 * GIB)];
+        let mut package = package(4, 18 * GIB);
+        package.layer_weight_bytes = vec![GIB / 8, GIB / 8, 9 * GIB, 8 * GIB];
+
+        let plan = plan_runtime_slice_topology_with_resources(
+            "topology-test",
+            "model-a",
+            &package,
+            &participants,
+            &[],
+            SplitTopologyResourceInputs {
+                native_context_length: 1,
+                kv_bytes_per_token: 1,
+                ctx_size_override: Some(1),
+                parallel_override: Some(1),
+            },
+        )
+        .expect("resource-aware topology with exact layer weights");
+
+        assert_eq!(
+            plan.stages
+                .iter()
+                .map(|stage| (stage.layer_start, stage.layer_end))
+                .collect::<Vec<_>>(),
+            vec![(0, 3), (3, 4)]
+        );
+        assert_eq!(plan.stages[0].parameter_bytes, 9 * GIB + GIB / 4);
+        assert_eq!(plan.stages[1].parameter_bytes, 8 * GIB);
+    }
+
     #[test]
     fn resource_planner_prefers_lower_tpot_stage_count_from_participant_rtt() {
         let participants = vec![