Plan split topology with exact layer weights

[i386]

Title

Plan split topology with exact layer weights

Original problem

Resource-aware split planning treated transformer layers as uniformly sized. That produces poor boundaries for models with uneven dense, routed-expert, shared-expert, or auxiliary layers, and could undercount shared endpoint tensors and KV cost across parallel lanes.

Diagnostics

Synthetic uneven-layer fixtures showed that average-layer planning can reject valid placements or choose boundaries that overload a later node. Package inspection also showed that embeddings, output tensors, and other shared bytes were not represented in a per-layer vector.

Fix

Carry ordered per-layer weight bytes from layer-package metadata into topology planning and choose contiguous boundaries from cumulative exact weights.

Account for shared model bytes at the first/final endpoints, multiply KV requirements by actual parallel lanes, and evaluate each later stage from its current boundary. Fall back to existing average-layer behavior when package indices are incomplete or non-contiguous.

The planner also accepts explicit context overrides below the automatic 64K planning floor while continuing to reject values above native context.

Validation

• cargo check -p skippy-coordinator -p mesh-llm-host-runtime -p mesh-llm
• cargo test -p skippy-coordinator (29 passed)
• cargo test -p mesh-llm-host-runtime inference::skippy::package::tests --lib (9 passed)
• cargo test -p mesh-llm-host-runtime runtime::split_planning::tests --lib (7 passed)
• cargo clippy -p skippy-coordinator -p mesh-llm-host-runtime -p mesh-llm --all-targets --no-deps -- -D warnings
• cargo fmt --all --check
• git diff --check
• No UI changes; screenshots do not apply.

Summary by CodeRabbit

• New Features

  • Improved model topology planning by accounting for the actual weight size of each layer.
  • Better supports uneven layer distribution across nodes based on available memory.
  • Added per-layer weight information to package and planning data.

• Bug Fixes

  • Corrected handling of shared model weight data when calculating layer sizes.
  • Context-length overrides below the planner’s minimum are now accepted when otherwise valid.

• Tests

  • Added coverage for uneven layer weights, stage boundaries, and invalid layer ordering.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change adds per-layer weight byte metadata to Skippy package identities, forwards it through runtime split planning, and uses it to calculate topology stage boundaries from uneven layer requirements. Explicit context overrides below the automatic minimum are now accepted.

Changes

Per-layer weight planning

Layer / File(s)	Summary
Package identity weight derivation crates/mesh-llm-host-runtime/src/inference/skippy/package.rs, crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs, crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs	SkippyPackageIdentity now stores per-layer weight bytes. HF layer metadata derives validated weights and distributes shared bytes across endpoint layers; direct GGUF and test identities use empty vectors.
Runtime planning input propagation crates/mesh-llm-host-runtime/src/runtime/split_planning.rs, crates/mesh-llm-host-runtime/src/runtime/local.rs	Runtime planning forwards package weights when their count matches the layer count, with updated fixtures and coverage for exact package weights.
Topology fitting and context overrides crates/skippy-coordinator/src/topology.rs	Topology fitting computes cumulative per-layer weight and KV requirements for stage boundaries, retaining uniform-weight fallback behavior. Explicit context overrides below the automatic floor are accepted and tested.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: ndizazzo, ivgolovach

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: split topology planning now uses exact per-layer weights.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch jd/exact-layer-weight-topology

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review.

[coderabbitai]

🧹 Nitpick comments (1)

crates/skippy-coordinator/src/topology.rs (1)

296-298: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Defend against division-by-zero panics on invalid input.

If input.layer_count is 0, div_ceil will panic. While upstream package resolution logic validates that layer counts are strictly positive, the TopologyPlanningInput contract itself does not strictly enforce this in the struct fields, making the coordinator potentially vulnerable to crashes from malformed inputs.

Consider clamping the divisor to a minimum of 1 to safely handle zero-layer inputs.

• crates/skippy-coordinator/src/topology.rs#L296-L298: Use u64::from(input.layer_count).max(1) as the divisor.
• crates/skippy-coordinator/src/topology.rs#L448-L450: Use u64::from(input.layer_count).max(1) as the divisor.

🛡️ Proposed fixes

For lines 296-298:

     let kv_per_layer = input
         .kv_bytes_per_token
-        .div_ceil(u64::from(input.layer_count));
+        .div_ceil(u64::from(input.layer_count).max(1));

For lines 448-450:

     let weight_per_layer = input
         .model_weight_bytes
-        .div_ceil(u64::from(input.layer_count));
+        .div_ceil(u64::from(input.layer_count).max(1));

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/skippy-coordinator/src/topology.rs` around lines 296 - 298, Prevent
division-by-zero in both kv_per_layer calculations by using
u64::from(input.layer_count).max(1) as the divisor. Update the calculations at
crates/skippy-coordinator/src/topology.rs lines 296-298 and 448-450; no other
behavior needs to change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/skippy-coordinator/src/topology.rs`:
- Around line 296-298: Prevent division-by-zero in both kv_per_layer
calculations by using u64::from(input.layer_count).max(1) as the divisor. Update
the calculations at crates/skippy-coordinator/src/topology.rs lines 296-298 and
448-450; no other behavior needs to change.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 2ee23b4e-247f-4b34-94b6-40d2dfa3fdf6

📥 Commits

Reviewing files that changed from the base of the PR and between 4f25060 and 9258adc.

📒 Files selected for processing (6)

• crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs
• crates/mesh-llm-host-runtime/src/inference/skippy/package.rs
• crates/mesh-llm-host-runtime/src/inference/skippy/resolver/test_support.rs
• crates/mesh-llm-host-runtime/src/runtime/local.rs
• crates/mesh-llm-host-runtime/src/runtime/split_planning.rs
• crates/skippy-coordinator/src/topology.rs