Repository navigation
feat: harden paged KV handoff deserialization boundary - #186
Merged
Merged
Conversation
Add the untrusted-network hardening carried over from the #125 distributed paged serde to the cross-node restore boundary, and fill the remaining single-node paged correctness gap. Deserialization hardening (src/distributed/kv_cache_serde): - CacheIngestLimits frame/block/slab caps. deserialize_cache_state now rejects oversized buffers and forged metadata lengths before parsing. - ExpectedBlockGeometry anchors every transferred slab to the decode model's real (block_size, n_kv_heads, head_dim, dtype) before the first write would capture pool geometry from the wire (write_block only validates subsequent writes against the captured geometry). - validate_paged_payload enforces table/contents layer consistency (no cross-layer aliasing, no missing or orphan contents) and rejects metadata-only or empty-table paged handoffs that a fresh decode pool cannot satisfy. restore_into_cache_pool_sequence_anchored applies it on the live path; the existing entry point keeps its signature with default limits and no anchor. Correctness: - trim_across_shared_blocks_releases_only_the_trimming_sequences_refs proves a trim spanning a shared block drops only the trimming sequence's refcount and leaves the sibling's prefix bit-identical (the trim/rewind-across-shared-blocks scenario; rewind is an alias of trim). Drive-by: correct the stale mlx_dtype_element_size(BFLOAT16) assertion. #125 fixed the function to return 2 bytes, but the kv_cache_serde suite was not run with that fix so the test kept asserting the old value 4. Part of #126.
inureyes
added
type:security
inureyes
added a commit
that referenced
this pull request
Oct 27, 2026
#2671) ## 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 #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 #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
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
Add the untrusted-network hardening carried over from the lablup#125 distributed paged serde to the cross-node restore boundary, and fill the remaining single-node paged correctness gap. Deserialization hardening (src/distributed/kv_cache_serde): - CacheIngestLimits frame/block/slab caps. deserialize_cache_state now rejects oversized buffers and forged metadata lengths before parsing. - ExpectedBlockGeometry anchors every transferred slab to the decode model's real (block_size, n_kv_heads, head_dim, dtype) before the first write would capture pool geometry from the wire (write_block only validates subsequent writes against the captured geometry). - validate_paged_payload enforces table/contents layer consistency (no cross-layer aliasing, no missing or orphan contents) and rejects metadata-only or empty-table paged handoffs that a fresh decode pool cannot satisfy. restore_into_cache_pool_sequence_anchored applies it on the live path; the existing entry point keeps its signature with default limits and no anchor. Correctness: - trim_across_shared_blocks_releases_only_the_trimming_sequences_refs proves a trim spanning a shared block drops only the trimming sequence's refcount and leaves the sibling's prefix bit-identical (the trim/rewind-across-shared-blocks scenario; rewind is an alias of trim). Drive-by: correct the stale mlx_dtype_element_size(BFLOAT16) assertion. lablup#125 fixed the function to return 2 bytes, but the kv_cache_serde suite was not run with that fix so the test kept asserting the old value 4. Part of lablup#126.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
First slice of the epic #116 capstone (#126). It hardens the cross-node paged
KV restore boundary with the untrusted-network items carried over from the #125
distributed paged serde, and fills the one remaining single-node paged
correctness gap. No live transport is wired here (that is the next slice); these
checks apply to the existing
deserialize_cache_state/restore_into_cache_pool_sequenceboundary, which the disaggregated handoff already feeds with untrusted bytes.
Deserialization hardening (
src/distributed/kv_cache_serde)CacheIngestLimits(frame / block-count / per-slab byte caps).deserialize_cache_staterejects oversized buffers and forged metadata lengthsbefore parsing. Defaults are generous backstops; the live decode path passes a
config-derived cap.
ExpectedBlockGeometrypins every transferred slab to thedecode model's real
(block_size, n_kv_heads, head_dim, dtype). This mattersbecause
PagedBlockPool::write_blockcaptures pool geometry from the firstslab it writes and only validates subsequent writes against it
(
cache/paged.rs:1006-1029), so without the anchor a peer could establish awrong-shaped pool on the first block.
validate_paged_payloadrejectscross-layer aliasing, block-table references with no transferred contents,
orphan contents (acquired but unreferenced, a pool-block leak), and duplicate
block ids disagreeing on layer.
table that references blocks with no contents (metadata-only), is rejected: a
fresh decode pool holds no pre-existing block contents to re-register.
restore_into_cache_pool_sequence_anchoredapplies the full set on the livepath; the existing
restore_into_cache_pool_sequencekeeps its signature withdefault limits and no model anchor (still enforces the structural checks).
Correctness gap-fill
trim_across_shared_blocks_releases_only_the_trimming_sequences_refs: twopool-backed sequences share a 2-block prefix (refcount 2); trimming one across
the shared tail block drops only that sequence's refcount (2 -> 1, never
freed), leaves the sibling's gathered prefix bit-identical, and shrinks the
trimmed sequence's own window to the retained prefix. Covers the
trim/rewind-across-shared-blocks scenario (rewind is an alias of trim).
Drive-by
mlx_dtype_element_size(BFLOAT16)returns 2 bytes (#125 fixed the function), butits assertion was left at the old value 4 because the
kv_cache_serdesuite wasnot run alongside that fix. Corrected.
Testing (orchestrator-local, Apple Silicon / Metal)
cargo test -p mlxcel --lib kv_cache_serde: 41 pass (18 new hardening tests).cargo test -p mlxcel-core --lib cache::paged: 75 pass (incl. the new trim test).cargo test -p mlxcel --test paged_kv_serialize_parity --release -- --ignored:qwen3-0.6b + llama-3.2-1b real-model handoff round-trip still byte-identical
(the new validation does not reject legitimate handoffs).
cargo cleanfirst, both crates,--all-targets,-D warnings)clean; rustfmt clean.
Part of #116.
Part of #126.