Adapt MLX pipeline concepts for Skippy

[i386]

Summary

Skippy can now borrow the useful parts of MLX pipeline parallelism without becoming an MLX backend and without adding tensor-parallel collectives.

The important translation is this:

• MLX pipeline mode is good at owning only the layers a rank needs, moving activations along a neighbor pipeline, and making request/rank state explicit.
• Skippy already has the right architectural shape for this with GGUF layer packages and stage servers.
• This PR makes Skippy’s existing pipeline more intentional: ordered transport edges, cold-start artifact accounting, explicit stage roles, request/cache epochs, and shared transfer-size accounting.

No public OpenAI API behavior changes. No stage wire-format bytes change. The new behavior is planner/protocol-helper/telemetry groundwork for better split serving.

Big Picture

Before, Skippy’s pipeline mostly answered: “which nodes can hold these layer ranges?”

After this PR, the pipeline can also answer: “which ordered chain is cheapest to run, what artifacts will it need, what role does each stage play, what cache epoch is this message in, and how large is the activation transfer we are timing?”

flowchart LR
  Resolve["Package metadata\nlayer inventory"] --> Plan["Topology planner\ncapacity + cache + edge cost"]
  Plan --> Roles["Stage roles\nDriver / Embedding / Intermediate / Readout"]
  Roles --> Runtime["Stage runtime\nrequest epoch + activation transfer telemetry"]
  Runtime --> Diagnose["Diagnostics\nwhy this plan, what it costs"]

Loading

1. Transport-Aware Stage Ordering

Before

The package-aware contiguous planner could prefer nodes based on node-local signals such as VRAM, cached artifact bytes, missing artifact bytes, and RTT-ish placement hints.

That was useful but incomplete for a pipeline. In a layer pipeline, adjacent stages exchange activations on every prefill/decode step. The cost that matters is not only “is node B good?” but also “is the edge from node A to node B good?”

So two plans could contain the same nodes and the same layer ranges, but have very different runtime behavior:

flowchart LR
  subgraph SameNodesBadOrder["Before: same selected nodes, unlucky order"]
    A["stage 0\nnode A"] -->|"18 ms"| B["stage 1\nnode B"]
    B -->|"24 ms"| C["stage 2\nnode C"]
    C -->|"3 ms"| D["stage 3\nnode D"]
  end

Loading

The old planner had no directed edge model to say that A -> C -> B -> D is a better chain than A -> B -> C -> D when both are otherwise valid.

How It Works Now

This PR adds StageEdgeSignal:

• source_node_id
• target_node_id
• optional directed rtt_ms
• optional large_frame_bytes_per_sec
• direct_prediction_return_supported

The new plan_package_aware_contiguous_with_transport(...) entry point keeps the same package-aware placement constraints, then scores possible stage orderings by adjacent directed edge cost.

For small stage counts, it exhaustively evaluates permutations. For larger stage counts, it uses a greedy fallback so the planner does not explode combinatorially.

flowchart LR
  subgraph BetterOrder["After: ordered by adjacent edge cost"]
    A["stage 0\nnode A"] -->|"2 ms"| C["stage 1\nnode C"]
    C -->|"3 ms"| B["stage 2\nnode B"]
    B -->|"4 ms"| D["stage 3\nnode D"]
  end

Loading

Diagnostics now record measured chosen edges, for example:

[info] pipeline edge mac-a -> mac-c after stage-0 has measured rtt 2 ms
[info] pipeline edge mac-c -> mac-b after stage-1 has measured rtt 3 ms

Why This Is Good

Skippy split serving is latency-sensitive because decode is a serialized chain. A single bad adjacent hop can dominate the whole pipeline even if every individual node has enough memory.

This gives the planner the same practical discipline MLX gets from pipeline ranks over Ring/TCP: make the model topology match the transport shape. It is especially useful on real meshes where Ethernet, Wi-Fi, Thunderbolt, and remote nodes can all appear in the same candidate set.

Reviewer Notes

• This does not use host-runtime peer RTT as if it were true stage-to-stage RTT.
• The API accepts directed edge signals explicitly, which keeps the planner honest until runtime stage-edge measurement is wired in.
• Existing package-aware planning still works via the old entry point.

2. Artifact Cold-Start Visibility

Before

Skippy already has layer packages, and package locality can affect placement. But once the planner accepted a topology, it did not summarize the artifact cold-start cost of that topology.

Operators could infer some things from node selection, but they could not easily see:

• how many selected bytes are already cached
• how many bytes are missing
• how many missing bytes can come from peers
• how many missing bytes must fall back to remote download

That made slow startup harder to explain.

How It Works Now

This PR adds artifact_diagnostics.

After a package-aware plan is selected, the planner aggregates selected-stage artifact data:

• cached_slice_bytes
• missing_artifact_bytes
• whether the selected node supports artifact transfer
• remote-download fallback bytes

Representative diagnostic:

[info] artifact cold-start plan: cached=19327352832 bytes, missing=6442450944 bytes, peer-transfer-eligible=4294967296 bytes, remote-download-fallback=2147483648 bytes

The diagnostic stays quiet for plain weighted plans that did not provide artifact signals, so existing non-package planning does not get noisy placeholder output.

Why This Is Good

This is the Skippy version of MLX’s metadata-first load trick.

MLX pipeline load first discovers metadata and rank-owned weights, then each rank downloads only the files it needs. Skippy cannot copy that mechanism directly because it serves GGUF layer packages rather than safetensors, but the idea transfers cleanly:

1. inspect package/layer metadata first
2. plan topology using cache and transfer signals
3. materialize only the accepted stage artifacts
4. explain the cold-start cost of that decision

This makes package planning less of a black box and gives us the data needed for future “prefer cached topology” and peer-transfer scheduling work.

3. Stage Role Metadata

Before

stage_index carried too much meaning implicitly.

For example:

• stage 0 usually acts as the driver
• the first layer owner acts as the embedding stage
• middle stages are intermediate activation processors
• the final layer owner acts as the readout/logits stage

Those roles were real, but they were not explicit in the topology data. A reviewer or UI had to infer behavior from index and layer range.

How It Works Now

This PR adds StageRole:

• Driver
• Embedding
• Intermediate
• Readout

Each StagePlan now carries a roles: Vec<StageRole> field with serde defaulting, so older serialized plans that do not include roles still deserialize.

Role assignment is deterministic:

• first stage: Driver, Embedding
• middle stages: Intermediate
• last stage: Readout
• single-stage plan: Driver, Embedding, Readout

flowchart LR
  S0["stage 0\nroles: Driver, Embedding\nlayers 0..12"] --> S1["stage 1\nroles: Intermediate\nlayers 12..28"]
  S1 --> S2["stage 2\nroles: Readout\nlayers 28..40"]
  S2 -. "prediction return" .-> S0

Loading

Why This Is Good

MLX’s rank ordering made final/readout ownership obvious by convention. Skippy should not reverse public stage numbering just to copy MLX, because stage 0 owns driver/session responsibilities.

Explicit roles give us the useful part without destabilizing Skippy’s mental model:

• diagnostics can say what a stage does, not just what number it has
• UI/API surfaces can show role-aware topology
• future direct-return/readout experiments can reason about readout separately from numeric stage index
• single-stage and split-stage behavior are represented uniformly

This also reduces accidental coupling between “stage index” and “stage responsibility.”

4. Request / Cache Epoch Coordination

Before

The binary stage protocol already carried the necessary identity fields:

• request_id
• session_id
• checkpoint_generation
• prompt_token_count
• decode_step

But code treated them as individual fields. There was no single value that meant “this stage message belongs to this request/cache generation.”

That makes future overlapping prefill, restore, trim, decode, and cache-reuse logic harder to audit.

How It Works Now

This PR adds StageRequestEpoch, derived from existing wire/header fields:

StageRequestEpoch {
  request_id,
  session_id,
  checkpoint_generation,
  prompt_token_count,
  decode_step,
}

StageWireMessage::request_epoch() centralizes extraction.

The epoch helper can compare whether two messages belong to the same request/session flow and whether one epoch is stale relative to another epoch.

Binary transport telemetry now includes checkpoint generation with the other epoch values:

skippy.checkpoint_generation=3
skippy.prompt_token_count=8192
skippy.decode_step=0

Why This Is Good

MLX rank 0 keeps ranks synchronized by broadcasting request state through collectives. Skippy does not need that exact broadcast pattern because Skippy has typed stage messages, but it needs the same discipline: every stage should agree about which request/session/cache generation a frame belongs to.

This PR does not start rejecting stale frames yet. It creates the safe vocabulary for that next step:

• cache restore/trim logic can compare epochs directly
• diagnostics can report exact request/cache generation
• tool-call loops with prefix reuse become easier to debug
• stale frame handling can be added without inventing a second identity model

Compatibility

This is not a wire-format change.

• no new encoded fields
• no stage state version bump
• no protobuf/gossip change
• no mixed-version binary framing break

The protocol crate changed, but the on-the-wire bytes did not.

5. Activation-Frame Transfer Accounting

Before

Skippy had two separate approximations of stage-message size:

• telemetry estimated llama_stage.message_wire_bytes
• downstream wire conditioning estimated transfer size for simulated bandwidth/delay

Those estimates could drift. The conditioned-wire path also undercounted parts of the message such as the full fixed header, position sidebands, sampling metadata, chat metadata, and raw import bytes.

For large prefill activation frames, that undercount makes simulated network pressure less trustworthy.

How It Works Now

This PR adds StageWireMessage::estimated_wire_bytes().

The estimate includes:

• fixed stage wire header
• sampling config bytes
• logit bias bytes
• chat sampling metadata length prefix and bytes
• token sideband bytes
• position sideband bytes
• state import raw bytes
• activation payload bytes

Both telemetry and wire conditioning now use the same helper.

Representative telemetry:

llama_stage.message_wire_bytes=52429612
skippy.activation_bytes=52428800
skippy.checkpoint_generation=3
skippy.prompt_token_count=8192
skippy.decode_step=0

Why This Is Good

MLX can stripe large payloads over multiple neighbor sockets. Skippy is not adding striping in this PR because that needs capability negotiation, chunk headers, reassembly buffers, backpressure, and mixed-version fallback.

But before adding striping, Skippy needs accurate accounting:

• telemetry should say how big the actual frame was
• bandwidth simulation should sleep based on the same size telemetry reports
• benchmark reports should use numbers that include sidebands and metadata
• future striping thresholds can be based on a shared transfer-size model

This is the low-risk foundation for later large prefill-frame striping.

CLI / Operator Output

No stable CLI command output changes in this PR.

The new user-visible text is planner and telemetry diagnostic material. Existing doctor/status/reporting surfaces can display it once they consume the new topology diagnostics.

Simulated representative output:

[info] pipeline edge mac-a -> mac-c after stage-0 has measured rtt 2 ms
[info] pipeline edge mac-c -> mac-b after stage-1 has measured rtt 3 ms
[info] artifact cold-start plan: cached=19327352832 bytes, missing=6442450944 bytes, peer-transfer-eligible=4294967296 bytes, remote-download-fallback=2147483648 bytes

llama_stage.message_wire_bytes=52429612
skippy.activation_bytes=52428800
skippy.checkpoint_generation=3
skippy.prompt_token_count=8192
skippy.decode_step=0

What This PR Does Not Do

• It does not add an MLX backend.
• It does not implement tensor parallelism.
• It does not add all-reduce collectives to Skippy decode.
• It does not change the OpenAI frontend ownership.
• It does not change the stage binary wire layout.
• It does not implement activation-frame striping yet; it adds the accounting needed to evaluate that safely.

Validation

• cargo fmt --all -- --check
• cargo test -p skippy-protocol --lib
• cargo test -p skippy-topology --lib
• cargo test -p skippy-server --lib
• cargo clippy -p skippy-protocol --all-targets -- -D warnings
• cargo clippy -p skippy-topology --all-targets -- -D warnings
• cargo clippy -p skippy-server --all-targets -- -D warnings
• cargo check -p mesh-llm
• cargo clippy -p mesh-llm --all-targets -- -D warnings

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f39e0ba6-0756-4838-b2b5-ae53d2721469

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/skippy-mlx-pipeline-concepts

---

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

[i386]

Superseded by concern-focused PRs: #814 transport-aware stage ordering, #815 artifact cold-start diagnostics, #816 stage role metadata, #817 request/cache epoch telemetry, and #818 stage wire byte accounting.