fix(distributed): bound tensor wire decode sizes and fix int16 decode #2691

Problem / Background

The security review of #2671 (issue #2654) found three pre-existing defects in how a stage decodes tensors received from a peer. Cluster frames are authenticated and, since #2684, AEAD-protected, so only an authenticated peer can reach these paths, but a compromised or buggy peer can make a stage allocate multi-GB buffers from a small frame, abort the process on a failed allocation, or silently misdecode a tensor. All three still hold on main at 54d63b0. Implement every item below in a single PR, run the local gate once at the end, and keep pushes to a minimum.

The base64 tensor field encoding from #2684 (src/distributed/pipeline/wire_bytes.rs) was checked and is not affected: base64 decoding (wire_bytes.rs:124-135) produces at most 3/4 of the input string length, which the frame cap already bounds (tcp_transport.rs:143, 256 MiB default), and the raw-sequence visitor caps its preallocation at 64 KiB (wire_bytes.rs:151-156). The bytes it yields flow into deserialize_wire_tensor (pipeline/wire_tensor.rs:42) and then deserialize_tensor, so items 1 and 2 cover that path. No change to wire_bytes.rs.

1. Quantized output sizing trusts an attacker-chosen shape (MEDIUM)

Current behavior. TensorHeader::num_elements (src/distributed/tensor_protocol.rs:246-259) folds the shape with checked_mul and then unwrap_or(u64::MAX), so an overflowing shape saturates instead of failing. deserialize_tensor (src/distributed/tensor_serialize.rs:173-186) casts that with as usize and passes it to dequantize_int8 / dequantize_int4 (src/distributed/tensor_quantize.rs:102-145, 215-265). Both return vec![0u8; num_elements * 2] whenever the payload is shorter than 8 bytes, has a zero group count or size, or is shorter than the shape implies (tensor_quantize.rs:104,111,121 and 217,224,235). A frame carrying a QUANTIZED int8 header with shape [1 << 31] and an 8-byte payload therefore returns Ok with a 4 GiB zero tensor; shape [1 << 40] aborts the process on allocation failure; a saturated u64::MAX overflows num_elements * 2 (panic in debug, wrap in release). The zero fill also turns a malformed payload into silently wrong data instead of an error. Unquantized payloads are never checked against the shape inside deserialize_tensor at all.

Fix.

• Change TensorHeader::num_elements to pub fn num_elements(&self) -> anyhow::Result<u64> that bails with shape product overflow for {shape:?} on overflow (empty shape still returns Ok(0)). Update the one caller (tensor_serialize.rs:174, converting with usize::try_from, never as) and the test at tensor_protocol_tests.rs:173.
• Change dequantize_int8 and dequantize_int4 to return anyhow::Result<Vec<u8>> and validate the payload layout exactly before allocating output: payload at least 8 bytes; group_size >= 1; num_groups == num_elements.div_ceil(group_size) (so num_elements == 0 requires num_groups == 0 and returns Ok(Vec::new())); and payload.len() == 8 + 2 * num_groups + data_len with data_len = num_elements (int8) or num_elements.div_ceil(2) (int4), all computed with checked arithmetic. Any mismatch returns Err naming the expected and actual lengths. Remove every zero-fill return. The output capacity num_elements * 2 then uses checked_mul, and is bounded by the payload the peer actually sent (at most 2x for int8, 4x for int4). Update the doc comment that currently says the function panics on malformed input.
• Propagate the Result at the other callers: kv_cache_transfer/quantized.rs:283-288 (already in a Result fn, add ?) and kv_cache_transfer/streamed.rs:329-338.
• At the end of deserialize_tensor, call the existing validate_data_size(original_dtype, &header.shape, &payload) (tensor_serialize.rs:214) so every decoded tensor, quantized or not, matches its shape exactly. Reuse that helper; do not write a second size check. serialize_tensor already enforces the same rule (tensor_serialize.rs:83), so every tensor the encoder produces still round-trips.

Acceptance criteria.

• deserialize_tensor on a QUANTIZED int8 header with shape [1 << 31] and an 8-byte payload returns Err without allocating the shape-sized buffer (new test deserialize_rejects_quantized_shape_larger_than_payload in tensor_serialize_tests.rs; on current code it returns Ok with 4 GiB of zeros, so the assertion fails).
• deserialize_tensor on a QUANTIZED header with shape [u64::MAX, 2] returns Err containing overflow and does not panic (new test deserialize_rejects_quantized_shape_overflow).
• deserialize_tensor on an unquantized float32 header with shape [4] and a 12-byte payload returns Err (new test deserialize_rejects_unquantized_size_mismatch).
• tensor_quantize_tests.rs covers dequantize_int8 and dequantize_int4 rejecting a payload that is short, has group_size == 0, or has a num_groups inconsistent with num_elements, and accepting the empty tensor; all existing round-trip tests (tensor_serialize_tests.rs:148,174,266, tensor_quantize_tests.rs, kv_cache_transfer/streamed_tests.rs) still pass.

2. LZ4 decompression trusts the size prefix (MEDIUM)

Current behavior. compress (src/distributed/tensor_compress.rs:35-41) writes [u64 LE original_len][u32 LE original_len][LZ4 block] (the inner u32 comes from lz4_flex::compress_prepend_size). decompress (tensor_compress.rs:47-56) reads the u64 into _original_len and ignores it, then calls lz4_flex::decompress_size_prepended, which in lz4_flex 0.11.6 (src/block/decompress_safe.rs:342-360) runs vec![0; uncompressed_size] on the peer-supplied u32 before decoding a byte. A 13-byte payload declaring 0xFFFF_FFFF makes the stage allocate 4 GiB per frame. Callers: deserialize_tensor (tensor_serialize.rs:167-170) and reassemble_layer_payload (kv_cache_transfer/streamed.rs:315-342).

Fix.

• Change the signature to pub fn decompress(data: &[u8], max_len: usize) -> anyhow::Result<Vec<u8>>. Before allocating: require data.len() >= 12; convert the u64 prefix with usize::try_from; require the inner u32 prefix to equal it (error LZ4 size prefix mismatch); require declared <= max_len (error decompressed size {declared} exceeds limit {max_len}); and require declared <= body.len().saturating_mul(LZ4_MAX_EXPANSION) with const LZ4_MAX_EXPANSION: usize = 256 (the LZ4 block format cannot expand beyond about 255:1, so this ties the allocation to the bytes actually received without rejecting valid data). Then allocate exactly declared bytes, decode with lz4_flex::block::decompress_into(body, &mut out), and require the written length to equal declared. Keep the on-wire format unchanged, so no protocol version bump.
• Add pub const MAX_DECODED_TENSOR_BYTES: usize = 1 << 32; to tensor_protocol.rs next to MAX_METADATA_LEN (tensor_protocol.rs:59). 4 GiB equals the u32 limit of the inner lz4_flex prefix, so nothing that decompresses today is rejected by it.
• deserialize_tensor passes min(expected, MAX_DECODED_TENSOR_BYTES), where expected is computed from the header with checked arithmetic: unquantized, num_elements * element_size (div_ceil(2) for Int4); quantized int8, 8 + 3 * num_elements; quantized int4, 8 + 2 * num_elements + num_elements.div_ceil(2) (upper bounds with group_size >= 1; item 1 then validates the exact layout). Put the two quantized bounds in tensor_quantize.rs as pub fn int8_payload_max_len(num_elements: usize) -> Option<usize> and int4_payload_max_len, and reuse them in reassemble_layer_payload, which passes the same bound derived from header.original_num_elements (None level: original_num_elements * 2).
• Update the two existing tests that call decompress without a limit (tensor_compress_tests.rs:74,79).

Acceptance criteria.

• decompress on [0xFFFF_FFFF as u64 LE][0xFFFF_FFFF as u32 LE][0x00] with max_len = 1 << 20 returns Err whose message contains exceeds (new test decompress_rejects_oversized_length_prefix_before_allocating in tensor_compress_tests.rs; today the call allocates 4 GiB and only then fails with a generic LZ4 error).
• decompress rejects a payload whose u64 and u32 prefixes disagree, and one whose declared size exceeds 256 * body_len (new tests).
• 64 MiB of zeros round-trips through compress and decompress(.., 64 << 20), proving the expansion bound does not reject maximally compressible data.
• deserialize_tensor on a COMPRESSED float32 header with shape [4] whose prefix declares 3 GiB returns Err (new test in tensor_serialize_tests.rs); the existing compressed round trips (tensor_serialize_tests.rs:124,266) pass.

3. from_bytes decodes 16-bit integers one byte per element (LOW)

Current behavior. The C++ bridge from_bytes (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:473-497) has typed cases for uint32, uint64, int32, int64, float32, float16 and bfloat16, and sends everything else through case 1: default: (mlx_cxx_bridge.cpp:493-495), which builds the array from a uint8_t* and casts each byte to the target dtype. For INT16 (code 6) and UINT16 (code 2) it reads only the first half of the buffer and produces wrong values; FLOAT64 (11) and COMPLEX64 (13) are misdecoded the same way. No over-read happens because validate_raw_tensor (src/distributed/kv_cache_serde/types.rs:480) checks the byte count first, but the tensor is silently wrong. Reachable from a peer: tensor_dtype_to_mlx (pipeline/wire_tensor.rs:83-95) maps TensorDtype::Int16 to 6, serialize_mlx_array emits Int16 for any int16 MLX array, and a KV-cache RawTensorData carries any dtype code that mlx_dtype_element_size accepts (types.rs:458-469).

Fix.

• In from_bytes, add case 2: return make_array_typed<uint16_t>(...) and case 6: return make_array_typed<int16_t>(...), using the existing make_array_typed helper (mlx_cxx_bridge.cpp:465). Make the fallthrough explicit as case 0: case 1: case 5: (bool, uint8, int8, all 1-byte) followed by default:, with a comment that only 1-byte dtypes may reach it.
• from_bytes cannot return an error through the cxx bridge (src/lib/mlxcel-core/src/lib.rs:166), so reject FLOAT64 and COMPLEX64 on the Rust side: validate_raw_tensor returns Err("dtype {code} is not supported for wire reconstruction") for codes 11 and 13. Leave mlx_dtype_element_size unchanged. No serialize path emits these codes for activations or KV caches.

Acceptance criteria.

• New test from_bytes_round_trips_16bit_integers_bit_exactly in src/lib/mlxcel-core/src/ffi_tests.rs, next to from_bytes_round_trips_16bit_floats_bit_exactly (ffi_tests.rs:1402), builds INT16 and UINT16 arrays (for example astype(from_slice_i32(&[-32768, -1, 0, 1, 255, 256, 32767, 7], &[2, 4]), INT16)), round-trips them through array_to_raw_bytes and from_bytes, and asserts identical dtype, shape and bytes. It fails on current code.
• New test in pipeline/wire_tensor_tests.rs: an int16 MLX array through serialize_mlx_array then deserialize_wire_tensor comes back byte-identical (the end-to-end wire path).
• New test in kv_cache_serde/tests.rs: validate_raw_tensor rejects dtype 11 and 13 with correctly sized data.

Scope

In scope: src/distributed/tensor_protocol.rs, tensor_serialize.rs, tensor_quantize.rs, tensor_compress.rs, kv_cache_transfer/streamed.rs, kv_cache_transfer/quantized.rs, kv_cache_serde/types.rs, src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, and the matching *_tests.rs files.

Out of scope: the wire format and cluster PROTOCOL_VERSION (unchanged), pipeline/wire_bytes.rs (not affected, see above), and frame-size limits in tcp_transport.rs.

Verification

Build outside the GPU lock first (a cargo build inside gpu-lock run leaks the lock to the sccache daemon), then run GPU tests under it. Every CUDA test run needs --test-threads=1.

cargo fmt --all -- --check
cargo clippy -p mlxcel --features cuda --lib --tests -- -D warnings
cargo test --profile test-fast --features cuda --lib --no-run
gpu-lock run --tag tensor-wire-bounds -- cargo test --profile test-fast --features cuda --lib distributed:: -- --test-threads=1
gpu-lock run --tag tensor-wire-bounds -- cargo test --profile test-fast -p mlxcel-core --features cuda from_bytes_round_trips -- --test-threads=1
cargo test --release --features cuda --test pipeline_remote_real_models --no-run
gpu-lock run --tag tensor-wire-bounds -- cargo test --release --features cuda --test pipeline_remote_real_models pipeline_remote_runtime_ -- --ignored --test-threads=1

A pass means every new test above passes, the full distributed:: lib suite passes, and the pipeline_remote_runtime_* real-model parity tests on GB10 pass as they do on main (fresh worktrees need the untracked models symlink). Record the parity test results in the PR.

References

Found in the security review of #2671 for #2654; base64 field encoding from #2684.