[inureyes]

Pipeline tensor decodes were bounded only by the sender's own header, so an authenticated peer could force about 4 GiB of allocation from a ~21 MiB frame and hand the executor any shape. Receivers now state what they expect, checked right after TensorHeader::decode and before the payload is sliced, copied, or allocated.

What changed

• Item 1. DimBound, ShapeBound, TensorDecodeLimits and deserialize_tensor_with_limits in tensor_serialize.rs (re-exported); deserialize_tensor wraps it unrestricted. Shape products use checked u64 arithmetic. deserialize_activation always refuses COMPRESSED and QUANTIZED. wire_tensor.rs adds MAX_STAGE_TOKENS_PER_FRAME (1 << 20), StageWireLimits and deserialize_wire_tensor_with_limits. RemoteStageService::load precomputes the hidden-state, input-id and mask limits once (PeerDecodeLimits) and uses them at the four peer-facing decode sites. RemotePipelineRuntime decodes logits under logits(input_seq_len) with its vocab_size, keeps validate_logits_shape, and now cleans up the request if that decode fails. reassemble_layer_payload takes max_layer_bytes and rejects a larger original_bytes before decompressing.
• Inbound width. StageExecutor/FamilyStageExecutor::inbound_activation_width default to None. LoadedStageExecutor stores the family override or loading::read_hidden_size (text_config first) and fails the load if both are absent. Gemma 4 (Gemma4 and Gemma4Vlm) overrides it through Gemma4StageModel::inbound_activation_width: hidden_size + (num_hidden_layers - layer_range.start) * hidden_size_per_layer_input when that field is non-zero, else hidden_size. unpack_stage_hidden slices with the same helper, so the formula exists once.
• Family audit. Every other adapter (llama, mistral, mixtral, qwen3, qwen35 incl. VLM, qwen3_next, gemma3, glm4 family, glm_moe_dsa, deepseek_v3, gpt_oss, jamba, llama4, nemotron_h) passes plain [1, L, hidden_size] across the boundary; none packs extras. VLM and Llama 4 configs nest hidden_size under text_config, which read_hidden_size reads first. All local checkpoints of these families resolve a width.
• Item 2. LayerTransferHeader.original_bytes; num_elements = raw_bytes.len() / 2; Int8/Int4 refuse an odd layer; level None bounds at original_bytes; Int8/Int4 require original_bytes == original_num_elements * 2.
• Item 3. Level None must decode to exactly original_bytes, compressed or not.
• Item 4. [] is a 1-element scalar via TensorHeader::shape_num_elements, shared with validate_data_size; the 18-space message is fixed.
• Item 5. Decode works from a borrowed Cow; only an unflagged payload is copied.
• Item 6. MAX_DECODED_TENSOR_BYTES: u64 plus max_decoded_tensor_bytes(). A 32-bit build cannot be run here (mlxcel-core does not build for 32-bit); the constant now const-evaluates there.
• Docs and existing tests. docs/distributed.md describes the receive-side bounds. forward_rejects_logits_with_wrong_vocab_dimension now expects the header-level vocab rejection, and validate_logits_shape gained a case for a sequence length strictly between 1 and the input length.

Failing-first evidence

Every new test was shown to fail by a mutation check on this branch (reverting the specific guard to main 752db4f's behavior), because each targets an API main lacks.

Test	Mutation	Failure
limits_refuse_flags_dtype_and_shape_before_the_payload, hidden_states_shape_is_checked_before_the_payload, entry_tensor_limits_bound_ids_and_masks, logits_limits_bound_sequence_and_vocab	drop check_header	"buffer too short" instead of the limit
activation_decode_refuses_compressed_and_quantized_frames	stop forcing flags off	LZ4 expansion error, no flag named
run_entry_with_out_of_bounds_ids_answers_limit_error, activation_outside_stage_limits_never_executes	unrestricted limits at the call sites	no limit error / executor ran
per_layer_input_inbound_width_matches_upstream_packing	width ignores per-layer inputs	256 != 304
read_hidden_size_reads_text_config_then_top_level	top level first	1152 != 2560
reassemble_rejects_original_bytes_above_max_layer_bytes	drop the bound	reached decompress
reassemble_round_trips_odd_byte_layer_compressed	per-tensor halves count	4095 != 4096
prepare_rejects_odd_byte_layer_for_quantization	drop odd check	Ok
reassemble_rejects_quantized_header_bytes_mismatch	drop bytes check	Ok, 512 bytes
reassemble_rejects_uncompressed_length_mismatch	drop length checks	Ok, 511 bytes
test_header_num_elements_empty_shape_is_one, float16_scalar_round_trips_and_empty_scalar_is_rejected	[] is 0 elements	0 != 1, serialize refuses scalar
decode_does_not_copy_the_payload_before_decompressing	to_vec() before decompress	assertion
test_max_decoded_tensor_bytes_helper	const back to usize	E0308 compile error

Verification

All on spark-102 (GB10) at aaa4daa, --profile test-fast, --test-threads=1, GPU runs under gpu-lock run, real-model runs with MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1. Every command exited 0.

Command	Result
cargo fmt --all -- --check	clean
cargo clippy -p mlxcel --features cuda,test-utils --lib --tests -- -D warnings	clean
cargo test -p mlxcel --features cuda --lib distributed::	1512 passed, 4 ignored
cargo test -p mlxcel-core --features cuda from_bytes	6 passed
--features cuda,test-utils --test distributed_integration	20 passed
--features cuda,test-utils --test pipeline_e2e	30 passed
--test pipeline_remote_real_models -- --ignored --skip glm_moe_dsa	10 passed (llama incl. RDMA/Thunderbolt, gemma3, gemma4, glm4, gpt_oss, qwen3, qwen3.5)
--test pipeline_server_remote_real_models -- --ignored --skip glm_moe_dsa	7 passed
--test pipeline_ci_multi_stage_real_models -- --skip glm_moe_dsa, with and without --ignored	2 + 2 passed, including pipeline_multi_stage_three_host_real_model_parity
--test disaggregated_e2e	51 passed
--test disaggregated_serving_e2e -- --include-ignored	2 passed
--test disaggregated_router_e2e -- --include-ignored	7 passed

The gemma4 runs use gemma-4-e2b-it-4bit (hidden_size_per_layer_input=256), so honest packed frames pass the per-layer width bound. glm_moe_dsa needs the 418.6 GB glm5-4bit checkpoint, which cannot load on a 128 GB box.

Closes #2716

Known gap after merging with #2720

On a non-entry stage, a forward activation refused by the new limits still leaves handle_activation through ? before report_activation_failure, so the coordinator waits out stage_timeout instead of getting an immediate Error (stage 0 answers at once). #2720 routes only execution and downstream-send failures to the coordinator. Routing every handle_activation failure, including these decode rejections, to the coordinator is folded into #2714, which rewrites the reply target in the same function; activation_outside_stage_limits_never_executes will then require the Error.

[inureyes]

Implementation Review Summary

Intent

Bound every peer-facing pipeline tensor decode by the receiver's expected flags, dtype, and shape before the payload is touched, and fix the five low-severity wire defects from #2716.

Findings Addressed

• forward_rejects_logits_with_wrong_vocab_dimension (runtime_correlation_tests.rs) still asserted the post-decode validate_logits_shape message and failed, because the logits limits now refuse the wrong vocab width on the header. It now pins the pre-decode rejection (a91c64a). (HIGH)

Remaining Items

• validate_logits_shape is now reachable only for a sequence length strictly between 1 and the input length, and no test covers that case. (LOW)
• The mask bound caps every axis at MAX_STAGE_TOKENS_PER_FRAME, so an explicit mask whose key axis exceeds 1,048,576 cached tokens would be refused. The spec mandates this bound and no current server path sends masks over the remote pipeline. (LOW)

Verification

• All stated requirements implemented (all six items, each check right after TensorHeader::decode)
• No placeholder/mock code remaining
• Integrated into project code flow (four remote-stage decode sites plus coordinator logits)
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes (edits in remote_service.rs and runtime.rs stay at the decode sites and limits plumbing)
• Tests pass: fmt, clippy -D warnings, and 20 focused lib modules under --test-threads=1, including distributed::pipeline::runtime. The full distributed:: module, real-model, and e2e suites are still pending.

[inureyes]

Security and performance review (head a91c64a)

Threat model: an authenticated but compromised cluster peer sends crafted frames. No CRITICAL, HIGH, or MEDIUM findings, so no fixes were committed.

Checked and clean:

• Every peer-facing decode in src/distributed/ and src/server/ goes through deserialize_wire_tensor_with_limits (the four RemoteStageService sites and the RemotePipelineRuntime logits). The remaining unrestricted deserialize_wire_tensor callers in stage_worker.rs read in-process ActivationReceiver channels, and TransportStageEndpoint::recv_activation has no production caller.
• check_header runs before any slice, copy, decompress, or allocation and does no arithmetic. The shape product is checked_mul in u64 and try_from to usize. Uniform mask shapes (rank 4, axes up to 2^20) overflow into an error, not a wrap. With COMPRESSED and QUANTIZED always refused on activations, the largest allocation is bounded by the received frame.
• Cow path: an unflagged payload is copied once (into_owned), the same as the old to_vec. The compressed and quantized paths read the borrowed slice.
• reassemble_layer_payload rejects original_bytes > max_layer_bytes first. For Int8/Int4 it requires original_num_elements * 2 == original_bytes before decompressing or dequantizing, so the output stays within max_layer_bytes.
• remote_service.rs and runtime.rs change only at the decode call sites and the limits plumbing.

Remaining LOW items (not fixed):

• StageWireLimits::hidden_states(), attention_mask(), and logits() build their TensorDecodeLimits per frame, which costs two or three small Vec allocations per token (remote_service.rs:482, :487, :606, :612; runtime.rs:625). Next to the payload copies this is negligible. If it ever matters, precompute the three static limits in RemoteStageService::load.
• DimBound::AtMost admits 0, so [1, 0, width] hidden states and [1, 0] input ids pass the limits and reach the executor. This behavior predates the PR and is outside fix(distributed): tighten tensor decode limits to the model shape #2716. A follow-up could add a minimum of 1 on the sequence axis.
• In validate_data_size, the Int4 branch uses num_elements.div_ceil(2) as usize, which can truncate on a 32-bit target. This also predates the PR, and the activation path rejects Int4 anyway.
• When a forward activation is rejected, ? still propagates out of handle_activation. Sibling PR update(distributed): transport and remote pipeline robustness follow-ups #2712 owns that error routing.

Verification: lib tests built with --profile test-fast --features cuda. Clippy -D warnings and cargo fmt --check are clean. Under gpu-lock with --test-threads=1: tensor_serialize 22, tensor_protocol 19, pipeline::wire_tensor 15, pipeline::activation_transfer 33, kv_cache_transfer 67, and pipeline::remote_service 11 tests passed.