update(distributed): remote pipeline residuals from #2671, #2680, #2684 #2694

Problem / Background

#2671, #2680 and #2684 (closing #2654, #2658 and #2656) left seven non-security residuals in the remote pipeline: one observability gap and two avoidable per-token costs in the frame path, a misleading TOML header and a missing abort in zero-config bring-up, and an unused mask plus a coverage gap in the Gemma4 stage. Each item below was re-checked on main @ 54d63b0 and still holds; none was dropped.

Implement all seven in a single PR, run the local gate once at the end, and keep pushes to a minimum (a Rust push costs a full CI cycle). Each item has its own acceptance criteria and its own test.

1. Log a tampered kind byte as a rejected frame

Current: ProtectedStream::read_frame (src/distributed/tcp_frame.rs:108-133) parses the kind with MessageKind::try_from(header[0])? at line 114, before the body is read and opened. A kind byte flipped to a valid value fails FrameOpener::open (the header is AEAD associated data) and surfaces as FrameRejected, which both callers log through the rate limiter (src/distributed/tcp_transport.rs:686-697, log_rejected_frame at src/distributed/stage_assignment.rs:686-697). A kind flipped to an invalid value (0, 5..=255) returns a plain unknown message kind error, and the connection closes with no log line at all.

Fix: on a protected stream, keep the raw kind byte, check len against max_frame_size as today, read the body, run opener.open(&header, &mut body), and only then call MessageKind::try_from. A tampered invalid kind therefore fails authentication and returns FrameRejected, logged by the existing call sites with no change to them. An unknown kind that does authenticate (a version-skewed member) returns a distinct error unknown message kind {n} from an authenticated peer (not FrameRejected). Plain streams keep parsing the kind first. Rejected: mapping any invalid kind to FrameRejected from the header alone, because the oversize check deliberately does not classify header-only refusals as auth failures (see protected_stream_refuses_oversize_header_before_body).

• Test protected_stream_rejects_tampered_invalid_kind_as_frame_rejected in src/distributed/tcp_transport_frame_tests.rs (same handshake pattern as protected_stream_refuses_oversize_header_before_body): seal a valid Control frame on the client, overwrite byte 0 with 0xFF, send it; read_frame returns an error that downcasts to FrameRejected. Fails on current code.
• Test on a plain stream: an invalid kind byte still errors and does not downcast to FrameRejected.

2. Seal in place instead of copying the payload in write_frame

Current: TcpTransport::send (src/distributed/tcp_transport.rs:307-309) builds the payload with serialize_message (src/distributed/tcp_frame.rs:137-165), then write_frame (tcp_frame.rs:86-101) allocates a second buffer and copies header plus payload into it before FrameSealer::seal (src/distributed/frame_protection.rs:316-342). For a per-token logits frame that is a 405 KB copy per send.

Fix: in tcp_frame.rs add pub(crate) fn encode_message_frame(msg: &TransportMessage) -> (MessageKind, Vec<u8>), which allocates FRAME_HEADER_LEN + payload + MAX_TAG_LEN capacity once, leaves the first FRAME_HEADER_LEN bytes for the header, and writes the payload after them. Add pub const MAX_TAG_LEN: usize = 32 to frame_protection.rs (the largest FrameSuite::tag_len). Add a sync helper fn finish_frame(sealer: Option<&mut FrameSealer>, kind: MessageKind, frame: &mut Vec<u8>) -> Result<()> that writes kind and BE length into the reserved header and seals in place, and ProtectedStream::write_prepared_frame(&mut self, kind, frame: Vec<u8>) that calls it and writes once. send uses the new pair. write_frame(kind, &[u8]) stays for borrowed payloads (send_stream, RPC, stage assignment) and becomes "copy into a prepared frame, then write_prepared_frame", so there is one sealing path. serialize_message becomes #[cfg(test)] (only tcp_transport_tests.rs uses it) and returns the payload slice of encode_message_frame. Out of scope: the JSON-to-Bytes copy inside TransportMessage::Control, which needs a Transport trait change.

• Test finish_frame_seals_without_reallocating: for each FrameSuite::ALL and for None, encode_message_frame of a 405 KB Control message, record frame.as_ptr(), call finish_frame, assert the pointer is unchanged and the length grew by exactly the suite's tag length.
• Test: a 405 KB Control message sent with write_prepared_frame over a protected and a plain loopback pair is returned byte-identical by read_frame and deserialize_message.

3. Send per-token logits once, from the final stage to the coordinator

Current: with two or more stages the final stage sends logits as a reverse ActivationMessage to its upstream (src/distributed/pipeline/remote_service.rs:681-702); each stage relays it upstream, and the entry stage turns it into RemoteStageResponse::RunResult for the coordinator using its PendingEntryReply map (remote_service.rs:164-169, insert at 527, relay at 576-597). So each decoded token carries the logits twice (about 405 KB base64 JSON each for qwen3) in 4 frames: RunEntry, forward activation, reverse activation, RunResult.

Fix:

• ActivationMessage (src/distributed/pipeline/activation_transfer.rs:46-71) gains #[serde(default)] pub reply_nonce: Option<u64> and pub fn with_reply_nonce(self, nonce: Option<u64>) -> Self. forward and reverse keep their signatures and set None.

• The entry stage (handle_run_entry, HiddenStates branch) sets reply_nonce = Some(nonce) from the RunEntry exchange and no longer records a pending reply. Middle stages copy activation.reply_nonce into the forward message they build (remote_service.rs:653-680).

• The final stage, on Logits, sends RunResult { stage_index: self.stage_index, nonce, request_id, micro_batch_id, logits } to self.coordinator.configured() through send_response. A forward activation without reply_nonce reaching the final stage is answered with nothing and logged at warn (it cannot be correlated).

• Remove PendingEntryReply, the pending_entry_replies map, the reverse-path relay in handle_activation, and the pending_entry_replies field of RemoteStageResponse::Ack and State (update runtime.rs stubs at 750-830, runtime_correlation_tests.rs, and tests/pipeline_remote_real_models.rs:232, which keeps its in_flight_requests == 0 check). remote_service drops a reverse activation with a warn. stage_worker.rs (in-process) keeps its reverse path.

• Coordinator: ExchangeMatch (src/distributed/pipeline/runtime.rs:234-262) gains result_peer: &StagePeer. In exchange (runtime.rs:389-466) it is stage_peers.last() for RunEntry and the command's peer otherwise. A RunResult must come from result_peer; Ack, State and Error still come from peer. A single-stage pipeline is unchanged (first stage is last).

• Logits stay base64 in the JSON RunResult (feat(distributed): protect every cluster frame with per-connection AEAD #2684). Rejected: a raw TensorData frame for logits, which carries no nonce or request id and would need a second correlation path in exchange for about 100 KB per token.

• PROTOCOL_VERSION stays 3: cluster_handshake.rs did not exist in v0.8.0-beta.1, so 3 has not shipped.

• Test in remote_service_tests.rs: a final stage (index 1 of 2, stub executor returning logits) given a forward activation with reply_nonce: Some(N) from its upstream sends RunResult { stage_index: 1, nonce: N } to the coordinator transport, and the upstream peer receives no activation. Fails on current code.

• Tests in runtime_correlation_tests.rs with two stage peers: a RunResult from the last stage completes forward; the same RunResult from stage 0 is dropped and the exchange times out.

• Serde round trip of reply_nonce in activation_transfer_tests.rs, including a message without the field decoding to None.

• Decode throughput measured once at the end on a quiet box with pipeline_remote_decode_bench (qwen3-0.6b-4bit, 2 stages, loopback), main vs branch interleaved over 5 rounds; the PR reports the median of medians for both (main after feat(distributed): protect every cluster frame with per-connection AEAD #2684 is about 133 tok/s) and the frame count per token (3 instead of 4). The branch must not be slower than main beyond run-to-run spread.

4. Peer TOML header names the stage assignment flow

Current: a --pp-peer writes the TOML it received through render_deterministic_toml (src/server/startup.rs:3039-3041), whose header (src/distributed/cluster_init.rs:301-302) says # Generated by `mlxcel-server --pp-auto`. and promises a byte-identical file on rerunning the bring-up command, which is the coordinator's command, not the peer's.

Fix: move the body emission of render_deterministic_toml into a private render_toml_body(cluster). Keep render_deterministic_toml (coordinator header unchanged; all existing callers and tests keep using it). Add pub fn render_peer_assignment_toml(cluster: &ClusterConfig) -> String with the header # Written by `mlxcel-server --pp-peer` from the stage assignment committed by its zero-config coordinator. followed by # The coordinator owns this plan; the next committed assignment overwrites this file. and a blank line, then the same body. Export it from src/distributed/mod.rs:112 and use it at startup.rs:3039.

• Test in cluster_init_tests.rs: the peer rendering starts with the new first line, contains no --pp-auto, parses with ClusterConfig::from_toml, and its text after the header equals the coordinator rendering after its header.

5. Abort committed peers when the commit phase fails

Current: push_stage_assignments (src/distributed/stage_assignment.rs:1245-1312) sends Commit to all accepted peers concurrently (commit_peer, 1209-1233). A peer that answers Committed immediately delivers its CommittedAssignment (handle_coordinator_connection, 834-873) and starts a stage, writing its TOML. If another peer fails the commit, the coordinator bails with AssignmentPushFailed and exits, leaving the committed peers running a stage for a cluster that will never form, with no message telling them.

Fix:

• AssignmentFrame (271-284) gains Activate and Abort { detail: String } (kind() returns "Activate" / "Abort"). Bump ASSIGNMENT_PROTOCOL_VERSION (line 84) to 2.

• Peer: after writing Committed, wait on the same stream with read_frame_unbounded under shared.commit_timeout. On Activate, deliver the assignment as today. On Abort, any other frame, close, or timeout: store false into shared.committed, log at info coordinator {bound} aborted the assignment for {node_id}: {detail}; waiting for another assignment (or the matching reason), and return so the listener keeps waiting, as the pre-commit discard paths already do. The peer never writes its TOML or starts a stage before Activate.

• Coordinator: keep each stream from the commit phase. If every peer committed, send Activate to all; an Activate write failure is reported in AssignmentPushFailed for that peer. If any commit failed, send Abort { detail } (the failure summary, truncated to 1 KiB) to every peer that committed, best effort within the remaining deadline (write errors logged at debug), then bail with AssignmentPushFailed as today.

• Document in the module docs the remaining window: a coordinator that dies between two Activate writes leaves earlier peers activated, the same as a coordinator crash after a successful bring-up. Rejected: activating on Committed and watching for Abort in the background, because the stage would already be loading weights and binding its transport.

• Test in stage_assignment_tests.rs: a scripted coordinator (raw write_frame/read_frame, as in peer_waits_for_commit_beyond_frame_io_timeout) sends Offer, Commit, then Abort; the peer delivers nothing, and a second full exchange ending in Activate on a new connection is delivered. Fails on current code.

• Test: push_stage_assignments with one real peer and one scripted peer that accepts and closes on Commit fails naming only the scripted peer, and the real peer receives Abort and does not return an assignment.

• Existing commit tests and tests/zero_config_cluster_init.rs pass with the Activate step.

6. Gemma4 stage without a sliding layer builds no sliding mask

Current: text_prefill_masks (src/models/gemma4.rs:4420-4443) builds create_sliding_window_prefill_mask whenever window < l <= DENSE_PREFILL_MASK_MAX_TOKENS (4096), even when sliding_attn is None. Gemma4StageModel::execute_hidden (gemma4.rs:5106-5121) passes None for a stage whose layers are all full_attention (for example a one-layer full-attention stage), so the [l, l] mask is allocated and never read.

Fix: return (None, None) at the top of text_prefill_masks when sliding_attn is None. The dense caller (gemma4.rs:3692) is unaffected because every Gemma4 text model has a sliding layer.

• Test in src/models/gemma4_stage_tests.rs: super::text_prefill_masks(None, h, &[Cache::Standard(KVCache::new())], WINDOW + 8, WINDOW) returns (None, None). Fails on current code.
• Test: a three-stage split 0..1, 1..2 (one full-attention layer), 2..LAYERS matches the dense logits bitwise on a WINDOW + 8 prefill.

7. Stage parity under a --max-kv-size trimmed cache

Current: gemma4_stage_tests.rs covers prefill, over-window prefill and continuation prefill, but no case where KVCache::trim_front has advanced live_start (#430). Under trim the global layers return live_len keys and must stay on the maskless causal path in both models; the only regression test is the dense-only first_cache_live_len_sizes_global_mask_from_live_window_after_trim (gemma4.rs:6849).

Fix: add two_stage_split_matches_dense_after_front_trim_bitwise: prefill 12 tokens in the dense model and in stages 0..1 + 1..LAYERS, call trim_front(4) on every Cache::Standard in both cache sets, then run a 3-token continuation prefill and one decode step, asserting bitwise equality with assert_bitwise_equal for both. This is coverage, so it may pass on current code; confirm locally that it fails when the global mask is sized from offset instead of live_len (not committed).

• The test exists, passes, and is shown in the PR to be sensitive to an offset-sized global mask.

Verification

Build once outside the lock (cargo test ... --no-run), then run under gpu-lock on GB10 (a build inside gpu-lock run leaks the flock to sccache). CUDA tests need --test-threads=1.

cargo fmt --check
cargo clippy --profile test-fast --features cuda,test-utils -p mlxcel --lib --tests -- -D warnings
gpu-lock run --tag pp-residuals -- cargo test --profile test-fast --features cuda,test-utils -p mlxcel --lib -- distributed:: models::gemma4 --test-threads=1
gpu-lock run --tag pp-residuals -- cargo test --profile test-fast --features cuda,test-utils -p mlxcel --test zero_config_cluster_init -- --include-ignored --test-threads=1
MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run --tag pp-residuals -- cargo test --profile test-fast --features cuda,test-utils -p mlxcel --test pipeline_remote_real_models -- --ignored --test-threads=1 pipeline_remote_runtime_
MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run --tag pp-residuals -- cargo test --profile test-fast --features cuda,test-utils -p mlxcel --test pipeline_server_remote_real_models -- --ignored --test-threads=1 pipeline_server_remote_
MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run --tag pp-residuals -- cargo test --profile test-fast --features cuda,test-utils -p mlxcel --test pipeline_ci_multi_stage_real_models -- --include-ignored --test-threads=1
# last, once, quiet box: main and branch interleaved, 5 rounds
gpu-lock run --tag pp-decode-bench -- cargo test --profile test-fast --features cuda -p mlxcel --test pipeline_remote_decode_bench -- --ignored --nocapture --test-threads=1

Pass: all green, including the three-stage pipeline_multi_stage_three_host_real_model_parity (logits now skip the middle stage), and the throughput table in the PR body.

Out of scope

A two-node link measurement for #2656 (Thunderbolt / 10-25 GbE), which needs two hosts. Execution errors from non-entry stages still travel upstream with nonce 0 and end in a coordinator timeout; unchanged here.

Refs #2654, #2656, #2658, #2671, #2680, #2684