fix(distributed): tighten tensor decode limits to the model shape #2716

Problem / Background

Residuals from PR #2701 (issue #2691), listed in section 7 of TECHNICAL_REPORTS/2701-bound-tensor-wire-decode-sizes-20261028.en.md. #2701 turned every peer-controlled size into a checked value, but the remaining bounds come from the sender's own header rather than from what the receiver expects. All six items were re-checked on main @ 752db4f and all hold. Every one needs an authenticated cluster member (#2650 / #2674 handshake, #2684 per-frame AEAD), so this is compromised-peer hardening, not an unauthenticated hole.

Severity: item 1 is medium (an authenticated peer can force about 8 GiB of allocation per ~21 MiB frame on a pipeline stage, and unbounded activation shapes reach the stage executor). Items 2 to 6 are low (latent correctness and portability defects with no amplification).

Implement all six in one PR, run the Verification block once at the end, and keep pushes to a minimum. Each item below needs its own test that fails on current code (or, where the API changes, a test whose scenario fails on current code).

1. Bound activation decode by the receiver's expected shape (medium)

Current behavior. The only production caller of deserialize_tensor (src/distributed/tensor_serialize.rs:152) is the pipeline activation path: deserialize_wire_tensor (src/distributed/pipeline/wire_tensor.rs:42) calls ActivationMessage::deserialize_activation (src/distributed/pipeline/activation_transfer.rs:157). Honest senders never set COMPRESSED or QUANTIZED (serialize_activation uses SerializeOptions::default()), but the decoder honors both. With a consistent header the limit at tensor_serialize.rs:181-184 is min(header bound, MAX_DECODED_TENSOR_BYTES) and decompress (src/distributed/tensor_compress.rs:67-108) accepts up to 256x the block, so a ~16 MiB LZ4 block (~21 MiB base64, under the 256 MiB TCP frame cap at src/distributed/tcp_transport.rs:159) allocates 4 GiB, and reconstruct_mlx_array -> from_bytes copies it again. Nothing checks rank, width, or sequence length before the tensor reaches the executor. Peer-facing call sites: src/distributed/pipeline/remote_service.rs:473 (entry input_ids), :475 and :593 (attention mask), :589 (hidden states), and src/distributed/pipeline/runtime.rs:620 (logits, checked by validate_logits_shape at runtime.rs:286 only after the full decode). src/distributed/pipeline/stage_worker.rs:153,320 is in-process.

Fix.

• tensor_serialize.rs: add pub enum DimBound { Exact(u64), AtMost(u64) }, pub enum ShapeBound { Axes(Vec<DimBound>), Uniform { rank: std::ops::RangeInclusive<usize>, max_dim: u64 } }, and pub struct TensorDecodeLimits { pub allow_compressed: bool, pub allow_quantized: bool, pub dtypes: Vec<TensorDtype>, pub shape: ShapeBound } (dtypes is the header dtype; empty accepts any) with TensorDecodeLimits::unrestricted(). Add pub fn deserialize_tensor_with_limits(buf: &[u8], limits: &TensorDecodeLimits) -> Result<(DeserializedTensor, usize)> that checks flags, dtype, and shape immediately after TensorHeader::decode, before the payload is sliced, copied, or allocated. deserialize_tensor(buf) becomes a wrapper with unrestricted(). Re-export beside deserialize_tensor in src/distributed/mod.rs:227. Errors name the flag, the dtype, or the axis and its bound.
• deserialize_activation always refuses COMPRESSED and QUANTIZED, which removes the 256:1 amplification from every activation decode, in-process included.
• wire_tensor.rs: add pub const MAX_STAGE_TOKENS_PER_FRAME: u64 = 1 << 20 (the same ceiling as MAX_HANDOFF_PROMPT_TOKENS, src/server/batch/scheduler/handoff.rs:340), pub struct StageWireLimits { pub activation_width: Option<u64>, pub vocab_size: Option<u64> } with hidden_states(), input_ids(), attention_mask(), logits(input_seq_len: u64) returning TensorDecodeLimits, and pub fn deserialize_wire_tensor_with_limits(wire_bytes: &[u8], limits: &TensorDecodeLimits) -> Result<UniquePtr<MlxArray>>. Hidden states: Axes[Exact(1), AtMost(MAX), Exact(width)], dtypes Float16 / BFloat16 / Float32. Input ids: Axes[Exact(1), AtMost(MAX)], Int32. Mask: Uniform { rank: 2..=4, max_dim: MAX }, Bool / Float16 / BFloat16 / Float32. Logits: Axes[Exact(1), AtMost(input_seq_len), Exact(vocab)], Float16 / BFloat16 / Float32. A None width or vocab becomes AtMost(u64::MAX) on that axis. The context window was rejected as the token bound because rope-scaled checkpoints legitimately exceed max_position_embeddings.
• Width comes from the executor, not config.json alone: Gemma 4 packs per-layer inputs into the activation (pack_hidden_for_downstream at src/models/gemma4.rs:5263, unpack_stage_hidden at :5280), so its inbound width is hidden_size + (num_hidden_layers - layer_range.start) * hidden_size_per_layer_input when that field is non-zero. Add fn inbound_activation_width(&self) -> Option<u64> { None } to StageExecutor and FamilyStageExecutor (src/distributed/pipeline/stage_executor/mod.rs:109,122). LoadedStageExecutor::load_with_adapter reads hidden_size through a new crate::loading::read_hidden_size(model_dir) -> Option<usize> next to read_vocab_size (src/loading/mod.rs:502; text_config first, then top level, like context_window_from_config), stores backend.inbound_activation_width().or(hidden_size), and fails the load if both are absent. The Gemma 4 adapter (stage_executor/gemma4.rs, used for both Gemma4 and Gemma4Vlm) overrides it through a pub(crate) helper on Gemma4StageModel beside unpack_stage_hidden, so the formula exists once.
• RemoteStageService::load (remote_service.rs:197) builds StageWireLimits once from executor.inbound_activation_width() and uses it at :473, :475, :589, :593. RemotePipelineRuntime passes logits(input_seq_len) with its vocab_size (runtime.rs:223) at :620 and keeps validate_logits_shape for the "input length or 1" rule.
• KV side: the production disaggregated handoff (src/distributed/disaggregated/handoff_impl.rs:209-260) uses kv_cache_serde with no compression and already anchors each slab to ExpectedBlockGeometry (src/distributed/kv_cache_serde/types.rs:571), so it is unaffected. The only KV path with LZ4 is reassemble_layer_payload (src/distributed/kv_cache_transfer/streamed.rs:315), which has no production caller. Add a max_layer_bytes: usize parameter and reject header.original_bytes > max_layer_bytes (field from item 2) before decompressing; its doc comment says a receiver derives it as 2 * n_kv_heads * head_dim * element_size * token_budget from ExpectedBlockGeometry.

Acceptance.

• A header-only buffer (data_len set, no payload bytes) claiming hidden states [1, 8, width + 1], [2, 8, width], or [1, MAX_STAGE_TOKENS_PER_FRAME + 1, width] is rejected by deserialize_wire_tensor_with_limits with the shape error, not "buffer too short", which proves the check precedes the payload.
• A COMPRESSED activation whose 13-byte payload declares 4 GiB under a consistent shape is rejected by deserialize_activation with an error naming the flag, before decompress runs; same for QUANTIZED.
• A unit test pins the Gemma 4 inbound width for a non-entry stage to the width pack_hidden_for_downstream emits upstream.
• reassemble_layer_payload rejects original_bytes > max_layer_bytes before decompressing.
• Remote real-model parity passes unchanged, so no honest frame is rejected.

2. Streamed KV element count (low)

prepare_layer_payload counts keys.len() / 2 + values.len() / 2 (streamed.rs:279), while level None bounds decompression at original_num_elements * 2 (streamed.rs:325): for a 1-byte dtype with odd key and value byte counts the bound is 2 bytes short and a valid layer is rejected. The same miscount breaks Int8 / Int4, because quantize_int8 / quantize_int4 (src/distributed/tensor_quantize.rs:74,227) count (k + v) / 2 over the concatenation, so the header disagrees with the quantized layout.

Fix. Add pub original_bytes: usize to LayerTransferHeader (src/distributed/kv_cache_transfer/mod.rs:185), filled from the original_bytes prepare_layer_payload already returns, at streamed.rs:222 and parallel.rs:123. Set num_elements = raw_bytes.len() / 2. Return an error for Int8 / Int4 when raw_bytes.len() is odd (the quantizers assume float16 and drop a trailing byte). In reassemble_layer_payload, level None bounds at original_bytes; Int8 / Int4 require original_bytes == original_num_elements * 2 before the existing payload bounds. No serde(default): the header has no production receiver.

• Level None, 1-byte layer, keys 4095 zero bytes, values 4097 zero bytes, compress = true (assert the wire was compressed): prepare -> reassemble round-trips. Current code rejects it with "exceeds limit".
• Int8 with an odd total byte count returns an error from prepare_layer_payload.

3. Uncompressed level None length unchecked (low)

reassemble_layer_payload returns wire_data.to_vec() (streamed.rs:333) without comparing it with the header. Fix: at level None require wire_data.len() == header.original_bytes (uncompressed) or a decompressed length equal to it (compressed), before copying.

• Uncompressed level None wire data one byte short and one byte long are both rejected.

4. Empty-shape semantics (low)

TensorHeader::num_elements (src/distributed/tensor_protocol.rs:258-261) and validate_data_size (tensor_serialize.rs:259-260) treat [] as 0 elements; validate_raw_tensor (types.rs:504) folds from 1, matching MLX scalars. Today [] with 0 bytes passes deserialize_tensor and fails only in validate_raw_tensor, and the encoder refuses a real scalar. Fix: adopt MLX semantics ([] is 1 element) in num_elements, and make validate_data_size reuse its product instead of a second fold. Update test_header_num_elements_empty_shape_is_zero (tensor_protocol_tests.rs:177). Rejecting empty shapes was considered and dropped: the wire path ends in validate_raw_tensor, which already reads them as scalars. Also fix the message at tensor_serialize.rs:285, which embeds 18 spaces because its line-continuation backslash is missing.

• A float16 scalar (shape [], 2 bytes) round-trips through serialize_tensor -> deserialize_tensor -> reconstruct_mlx_array; [] with 0 bytes is rejected by deserialize_tensor.

5. Copy before decompress (low)

tensor_serialize.rs:169 copies the payload before decompressing or dequantizing. Fix: keep &buf[data_start..data_end], decompress or dequantize from the borrowed slice, and call to_vec() only when neither flag is set.

• No to_vec() precedes decompress or dequantize_*; existing compressed and quantized round-trip tests pass.

6. 32-bit portability (low)

pub const MAX_DECODED_TENSOR_BYTES: usize = 1 << 32 (tensor_protocol.rs:66) is a const-evaluation overflow on a 32-bit usize. Fix: make it u64 and add pub fn max_decoded_tensor_bytes() -> usize { usize::try_from(MAX_DECODED_TENSOR_BYTES).unwrap_or(usize::MAX) }, used at tensor_serialize.rs:182 and streamed.rs:330.

• A unit test pins the helper. A 32-bit build cannot be run (mlxcel-core does not build there); say so in the PR.

Out of scope

• Production handoff passes CacheIngestLimits::default() (16 GiB frame) at handoff.rs:192,436, although types.rs:530 says the live path derives tight values. No amplification; separate follow-up.
• LayerTransferHeader has no compressed flag, so a future streamed receiver cannot choose compressed; settle when that receiver is wired in.
• FLOAT64 / COMPLEX64 inside from_bytes without validate_raw_tensor (fix(distributed): bound tensor wire decode sizes and fix int16 decode #2701 report, section 7).

Verification

Build every target first with --no-run outside gpu-lock (a cargo build inside gpu-lock run hands the lock to the sccache daemon), then run the tests under it. Real-model runs set MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1.

cargo fmt --all -- --check
cargo clippy -p mlxcel --profile test-fast --features cuda,test-utils --lib --tests -- -D warnings
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --lib distributed:: -- --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel-core --profile test-fast --features cuda from_bytes -- --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda,test-utils --test distributed_integration --test pipeline_e2e -- --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --test pipeline_remote_real_models -- --ignored --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --test pipeline_server_remote_real_models -- --ignored --skip glm_moe_dsa --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --test pipeline_ci_multi_stage_real_models -- --ignored --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --test disaggregated_e2e -- --test-threads=1
gpu-lock run --tag decode-limits -- cargo test -p mlxcel --profile test-fast --features cuda --test disaggregated_serving_e2e --test disaggregated_router_e2e -- --ignored --test-threads=1

A pass is zero failures in each, with the real-model parity tests run, not skipped.

Related: #2654, #2671, #2691, #2701.