[inureyes]

Root cause

Two independent defects, found by comparing hidden states at the split boundary on GB10 before changing anything:

• Wire decoding widened float16 to float32. deserialize_wire_tensor rebuilt float16 activations with from_bytes_f16, a checkpoint loader that converts float16 to float32 everywhere except Apple. Checkpoints that load as float16 on CUDA (llama-3.2-1b, gemma-3-1b) then ran the downstream stage and LM head in float32. On llama-3.2-1b-4bit, stage 0 output was byte-identical to the dense hidden state at the split, in-process stage logits were byte-identical to dense, and only the wire round trip moved the logits (max abs diff 1.5e-2, float32 instead of float16). qwen3 passed because it stays bfloat16, which that loader keeps native.
• Gemma stage models built different masks. Gemma3StageModel and Gemma4StageModel built explicit causal and sliding masks for every prefill, while the dense forwards leave both None inside the window and take the maskless causal SDPA. In-process stage logits (no wire) differed from dense by 9.4e-2 (gemma-3-1b) and 3.4e-1 (gemma-4-e2b, split 13).

Fix

• Wire tensors are rebuilt in the dtype they were sent in, through the guarded kv_cache_serde::reconstruct_mlx_array, whose existing validate_raw_tensor check (from feat: harden paged KV handoff deserialization boundary #186) rejects a payload whose length does not match the shape before the FFI copy. The wire path previously bypassed that check, so negative dims, overflowing element counts, and a short payload were unchecked there (a heap over-read from an authenticated peer; this predates the PR).
• Gemma 3 (text_prefill_sliding_mask) and Gemma 4 (text_prefill_masks) mask decisions are shared helpers used by the dense model, the stage model, and Gemma3Backbone. Dense behavior is unchanged in every branch, and stages now also keep the 4096-token no-dense-mask cutoff from fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672.

All three families now give byte-identical logits to dense, in-process and over the wire. No tolerance was changed.

Regression coverage

Non-ignored lib tests: wire_tensor::tests (f16 round trip, plus short f16/f32 and empty-scalar payload rejection) and gemma3_stage_tests / gemma4_stage_tests, which split a synthetic model (stage 1 opening on a global layer) and require dense logits bit for bit for 2- and 8-token prefill, continuation prefill on a non-empty cache inside and across the sliding window, and decode. The f16 round trip and both 2-token split tests fail against main's sources (diffs 8.8e-3 and 2.3e-2).

GB10 CI coverage decision

The real-model remote parity tests are not added to CI: the PP workflow's two-host job runs on ubuntu-latest (CPU) with only the qwen3 fixture, and no GB10 job runs the PP suites or holds llama/gemma checkpoints. Instead run-pp-two-host-logical.sh runs the two weight-free split suites (its distributed::pipeline:: filter already covers the wire tests), and the workflow's path filter now includes the Gemma model and test files so a dense-mask change triggers them.

Validation

On GB10 (spark-102) with MLXCEL_REQUIRE_MODELS=1 MLXCEL_SDPA_DETERMINISTIC=1 --test-threads=1 at the final head: pipeline_remote_real_models 6/6 (llama, gemma3, gemma4, qwen3, qwen3.5 parity, llama drain), pipeline_server_remote_real_models 5/5, and pipeline_multi_stage_two_host_logical_smoke; pipeline_stage_executor_real_models 10/10 at the first commit (its gemma variants and llama worker loop also failed on main). Lib distributed::pipeline::, models::gemma3, models::gemma4, models::gemma3_backbone, kv_cache_serde, fmt and clippy pass. The new lib tests have not yet run on a CPU build; PP CI on this PR is the first. Metal not checked; on Apple the float16 wire path was already native, so only the mask change applies there.

Closes #2654

[inureyes]

Implementation Review Summary

Intent

Make the 2-stage remote pipeline byte-identical to the dense model for llama, gemma3 and gemma4 on GB10 by keeping wire activations in their sent dtype and giving the Gemma stage models the dense prefill masks.

Findings

No CRITICAL or HIGH findings; no fix commit.

• Dense paths unchanged: the Gemma 3 dense, Gemma3Backbone and Gemma 4 text-only branches are boolean-equivalent to main (the merged Gemma 4 long-prefill branch is still reached only for l > 4096 with no padding or overlay); external-mask, padded, batched-decode and verify branches are untouched.
• Wire: from_bytes with ids 9/12 reinterprets 2-byte halves on every backend; Metal f16 and native bf16 were already lossless, and under MLX_BF16_NATIVE=0 weights load as f32, so activations stay consistent with dense.

Remaining Items

• No continuation (non-empty cache) multi-token prefill case in the stage parity tests (MEDIUM)
• pipeline-parallel-ci.yml path filters omit src/models/gemma3*.rs and gemma4*.rs, and main CI never runs mlxcel lib tests on Linux, so a dense-mask change that skips PP paths will not run the new parity suites (MEDIUM)
• gemma3_stage_tests.rs header says the fixture is bf16; it is float16 (LOW)
• The binary assets failure comes from the FLUX.2 fixtures already on main (feat(imagegen): port FLUX.2 Klein text-to-image and editing #2660), not this PR

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass: on GB10, models::gemma3::, models::gemma4::, distributed::pipeline:: pass with and without MLXCEL_SDPA_DETERMINISTIC=1, and the three new regression tests fail when the fix sources are reverted to main (8.8e-3 and 2.3e-2 logit diffs, f16 wire dtype); clippy -D warnings and fmt clean

[inureyes]

Security and performance review

Fixed (HIGH), 25e93de: deserialize_wire_tensor passed a peer-supplied shape and payload to mlxcel_core::from_bytes, which copies product(shape) elements from a raw pointer without checking the payload length. A short payload with a large shape read past the heap buffer. The bug predates this PR (the old from_bytes_f16 branch did not check either), but this PR routes every dtype through from_bytes. The function now goes through kv_cache_serde::reconstruct_mlx_array, which checks for non-negative dims, computes the element product with checked arithmetic (an empty shape is a 1-element scalar), and requires the exact byte length. New tests reject short f16 and f32 payloads and an empty scalar payload. wire_tensor:: (8 tests) and kv_cache_serde:: (41 tests) pass on GB10, and clippy and fmt are clean.

Performance: no regressions. The stage prefill used to build a dense causal global mask and a sliding mask on every prefill, including above 4096 tokens. It now uses the dense path's helpers: global layers take maskless causal SDPA, the Gemma 4 path keeps the DENSE_PREFILL_MASK_MAX_TOKENS cutoff (#672), and the sliding band is built only past the window when the native route is not planned. This means less allocation per prefill. The added pipeline_hint calls do nothing unless MLXCEL_PIPELINE_GRANULARITY is set.

Remaining (pre-existing, outside this diff, authenticated peers only):

• MEDIUM: tensor_serialize::deserialize_tensor handles quantized payloads by sizing dequantize_int8/int4 output from header.num_elements(). That value saturates to u64::MAX on overflow and is multiplied by 2 unchecked, so a crafted header causes a panic or an allocation abort.
• MEDIUM: tensor_compress::decompress trusts the LZ4 size prefix (up to 4 GiB) with no cap tied to the shape.
• LOW: from_bytes sends int16 through the uint8_t default branch, so an int16 wire tensor would be decoded incorrectly. Activations never use int16.