fix(models): size dense prefill masks from live_len under trim (#419) by inureyes · Pull Request #420 · lablup/mlxcel

inureyes · 2026-06-24T09:12:46Z

Summary

The dense-KVCache sliding-window models (cohere2, gemma3n, olmo3) built their multi-token (l > 1) prefill attention masks from the cache's monotonic offset, for both the global mask (create_causal_mask) and the sliding mask (create_sliding_window_prefill_mask_dense). Under the server --max-kv-size path that assumption is wrong after a trim, and the mismatch aborts the whole server process. This sizes the masks from cache.live_len() instead, which is byte-identical on the common untrimmed path.

The bug

enforce_max_kv_size_for (called in continue_chunked_prefill, among others) calls KVCache::trim_front, which physically slices the buffer to the live window and advances live_start while offset keeps growing monotonically (to preserve RoPE relative positions). update_and_fetch (Fp16 and Int8) then returns only the live window of live_len = offset - live_start keys, not offset keys. On a multi-token continuation chunk after at least one trim (live_start > 0), a mask sized from the monotonic offset is WIDER than the returned K/V, so the mask key axis and the K/V key axis disagree. MLX broadcast_shapes throws, the cxx bridge turns the C++ exception into std::terminate, and the server process ABORTS (not a recoverable per-request error).

Reproduced on gemma3n-e2b-4bit with --max-kv-size 1024 --prefill-chunk-size 512 on a 2631-token prompt: crash on chunk 4 with [broadcast_shapes] Shapes (512,2048) and (1,8,512,1536) cannot be broadcast.

What changed

src/models/cohere2.rs (forward_impl and forward_with_embeddings_impl): build the global and sliding prefill masks from caches[..].live_len() instead of caches[..].offset.
src/models/gemma3n.rs (both forward variants): same change for global_mask / sliding_mask; the comment notes consistency with the KV-shared attention layer, which already slices its K/V to cache.live_len().
src/models/olmo3.rs (forward_impl): same change for the global and sliding masks.
src/lib/mlxcel-core/src/cache.rs: new cache-level regression test kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim.

These offset reads were mask-only. RoPE keeps using each layer's own monotonic cache.offset (read internally by the attention layer); no positional / layer-forward path changed. create_causal_mask, create_causal_mask_with_window_full, create_sliding_window_prefill_mask_dense, trim_front, update_and_fetch, and the RotatingKVCache consumers are untouched.

Why it is correct (and byte-identical when untrimmed)

For a query row q at monotonic abs position offset + q, the live key column k is at abs position live_start + k. The causal bound live_start + k <= offset + q reduces to k <= q + (offset - live_start) = q + live_len, which is exactly create_causal_mask(l, live_len); the window bound reduces the same way. The mask key dim l + live_len then matches update_and_fetch's returned l + live_len. When no trim has happened, live_start == 0 so live_len == offset and the masks are bit-for-bit the same, so greedy output on the common path is unchanged.

Test plan

cargo test --release -p mlxcel-core --features metal,accelerate cache:: (402 passed, incl. the new kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim)
cargo check --lib --features metal,accelerate
cargo clippy --lib --tests --features metal,accelerate -- -D warnings (clean)
cargo fmt --check

Closes #419

The dense-KVCache sliding-window models (cohere2, gemma3n, olmo3) built their multi-token (l > 1) prefill attention masks from the cache's monotonic offset, for both the global mask (create_causal_mask) and the sliding mask (create_sliding_window_prefill_mask_dense). That is correct only while no trim has happened. Under the server --max-kv-size path, enforce_max_kv_size_for calls KVCache::trim_front, which physically slices the buffer to the live window and advances live_start while offset keeps growing monotonically (to preserve RoPE relative positions). update_and_fetch (Fp16 and Int8) then returns only the live window of live_len = offset - live_start keys, not offset keys. On a multi-token continuation chunk after at least one trim (live_start > 0), a mask sized from offset is wider than the returned K/V, so the mask key axis and the K/V key axis disagree, MLX broadcast_shapes throws, the cxx bridge turns the C++ exception into std::terminate, and the server process aborts. Reproduced on gemma3n-e2b-4bit with --max-kv-size 1024 --prefill-chunk-size 512 on a 2631-token prompt: crash on chunk 4 with [broadcast_shapes] Shapes (512,2048) and (1,8,512,1536) cannot be broadcast. Fix: build the model-level prefill masks from cache.live_len() instead of cache.offset, for both the global and sliding masks. For a query row q at monotonic abs position offset + q, the live key column k is at abs position live_start + k, so the causal bound live_start + k <= offset + q reduces to k <= q + (offset - live_start) = q + live_len, exactly create_causal_mask(l, live_len); the window bound reduces the same way. The mask key dim l + live_len then matches update_and_fetch's returned l + live_len. When untrimmed, live_start == 0 so live_len == offset and the masks are byte-identical, keeping greedy output unchanged on the common path. These offset reads are mask-only: RoPE keeps using each layer's own monotonic cache.offset, which the attention layer reads internally. gemma3n's KV-shared attention layer already slices its K/V to cache.live_len(); aligning the model-level mask to the same bound makes the two consistent. Adds a cache-level regression test (kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim) that trims an Fp16 KVCache so live_start > 0, then asserts a continuation update returns live_len + m keys and that create_causal_mask(m, live_len) matches that key axis while create_causal_mask(m, offset) is strictly wider (the bug).

inureyes · 2026-06-24T09:24:30Z

Real-model validation across all three dense sliding-window models (before/after)

The bug is a full server-process abort, reproduced on all three affected models with --max-kv-size 1024 --prefill-chunk-size 512 and a ~2.6k-token prompt: chunked prefill drives a trim_front (live window capped at 1024), and the following continuation chunk builds an offset-sized mask wider than the trimmed K/V, so MLX broadcast_shapes throws and cxx hits std::terminate (Abort trap: 6).

Before (main 22c9694e3, the bug is preexisting) vs after (this PR, masks sized from live_len()):

Model	window	before (`--max-kv-size 1024`)	after
gemma3n-e2b-4bit	512	abort: `broadcast_shapes (512,2048) vs (1,8,512,1536)`	HTTP 200, coherent summary
command-r7b-4bit (cohere2)	4096	abort: `(512,2048) vs (1,32,512,1536)`	HTTP 200, no crash
olmo3-32b-4bit	4096	abort: `(512,2048) vs (1,40,512,1536)`	HTTP 200, coherent summary

The mask shape (512, 2048) is [l, l + offset] (l=512, offset=1536); after the fix it is [l, l + live_len] = (512, 1536), matching the live K/V key axis. Identical for all three (the shared failure mode), differing only in head count.

Semantic-correctness control: cohere2 with the fix binary and NO trim (--max-kv-size 8192, the full 2632-token prompt retained, the untrimmed path that is byte-identical to main) returns a coherent summary ("The history of computing is a long chain of incremental ideas, each building on the last..."). The degenerate text cohere2 emits at --max-kv-size 1024 is the expected consequence of capping its context to 1024 tokens, far below its 4096 sliding window (so it loses ~1600 of 2632 prompt tokens); the fix makes that opted-into truncation run correctly instead of crashing. gemma3n (window 512 <= 1024) and olmo3 stay coherent under the same trim.

inureyes · 2026-06-24T09:56:49Z

Security / performance review

Reviewed the offset -> live_len() change on cohere2, gemma3n, olmo3.

In scope (the 3 models): correct and complete, no findings

DoS / abort path closed. Both the global (create_causal_mask) and sliding (create_sliding_window_prefill_mask_dense) prefill masks now size from live_len(), which matches the live_len + l key axis that update_and_fetch returns after a trim_front. The helper produces [l, l + live_len], so the mask key axis equals the returned K/V key axis. The only remaining cache.offset reads on these models' prefill path are RoPE position inputs (correct: positional, not a shape), and the K/V comes from update_and_fetch (the live window). No other mask or K/V slice on these paths can exceed the live K/V, so a client cannot induce the broadcast_shapes abort on these 3 models under --max-kv-size after this change. The --max-kv-size trim loop trims every per-sequence cache against its own live_len() in lockstep, so the representative-cache live_len() used to build the shared global/sliding masks is valid for every layer of that class.
Performance: no regression. live_len() = offset - live_start is an #[inline] field subtraction, not a recompute. live_len() <= offset, so the mask is the same size or smaller than before, with no extra allocations or graph nodes. The change is confined to the l > 1 prefill branch; decode (l == 1) is untouched.
Integer arithmetic. trim_front clamps live_start <= offset (at most live_start == offset, giving live_len() == 0), so live_len() cannot underflow.
No untrusted-input parsing, I/O, or auth surface is touched. The regression test pins the cache invariant the fix relies on.

Out of scope, pre-existing (HIGH, recommend follow-up)

The same crash signature (a model-level prefill mask sized from a dense KVCache's monotonic offset, passed to a non-rebuilding SDPA path, reachable under --max-kv-size chunked prefill with a multi-token continuation chunk after a trim) still exists in other dense models. Confirmed by inspection:

nemotron_h: create_causal_mask(seq_len, caches[idx].offset())
mistral4: create_causal_mask(l, caches[0].offset)
jamba: create_causal_mask(seq_len, c[attn_idx].offset())

For contrast, glm4_moe_lite and minicpm3 already build this mask from seq_len() (= live window) and are safe, which shows the migration is partial. These models are not part of #419; a separate issue should apply the same offset -> live_len()/seq_len() change across the remaining dense models, validated per family.

Not merging.

inureyes · 2026-06-24T10:00:12Z

PR Finalization Complete

Summary

Tests: Adequate. The cache-level regression test kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim precisely pins the invariant this fix relies on. The 3-model E2E validation (abort to HTTP 200) was already recorded in the PR. No additional tests added; no gaps found.
Documentation: No doc update warranted. The --max-kv-size flag and the live_len / trim_front mask-sizing behavior are not documented in any file under docs/. The relevant explanation lives in the inline code comments, which were updated in all three model files and in the new test.
Code quality: No debug artifacts, no leftover TODOs. Comments are accurate and reference the correct issue numbers (fix: harden sliding-window >window prefill across non-Gemma models; guard causal_attention windowed K/V slice #408, fix(models): dense-cache sliding-window models still degenerate on rolled-over over-window prefill (offset>0 && seq>window) #413, fix(server): dense sliding-window models build prefill masks from monotonic offset, not live_len, under --max-kv-size #419). All checks pass: cargo fmt --check (clean), cargo check --lib (clean), cargo clippy --lib --tests -- -D warnings (clean), cargo test --release -p mlxcel-core cache:: (402 passed, 0 failed).

All checks passing. Ready for merge.

…tron_nas/qwen-vl under --max-kv-size (#421) Completes the offset->live_len prefill-mask migration started in #419/#420 for the remaining dense Vec<KVCache> scheduler-pool models that build their multi-token prefill mask from the cache's monotonic offset. Under the server --max-kv-size path, enforce_max_kv_size_for calls KVCache::trim_front, which advances live_start while offset keeps growing for RoPE. update_and_fetch then returns only live_len = offset - live_start keys, so a mask sized from offset is wider than the live K/V, MLX broadcast_shapes throws, and the non-Result C++ exception crosses the cxx boundary to std::terminate and aborts the whole server. A remote client can crash all sessions with one long prompt plus a multi-token continuation chunk. Fix is mask-only: build the prefill causal mask from cache.live_len() (alias of seq_len(), both = offset - live_start) instead of offset. RoPE, position_ids, and the Llama-4 attention scale keep reading the monotonic offset. Byte-identical when untrimmed (live_start == 0, so live_len == offset). Models changed: - mistral4: transformer_body mask uses caches[0].live_len(); the Llama-4 attention scale below keeps caches[0].offset. - nemotron_nas: forward no-mask branch uses caches.first().map(|c| c.live_len()); per-layer RoPE offset untouched. - qwen2_vl / qwen3_vl / qwen3_vl_moe: only the auto_mask create_causal_mask argument switches to caches[0].live_len(); cache_offset stays monotonic for position_ids/RoPE grid slicing. nemotron_h and jamba were investigated and excluded: they are ModelOwnedSequenceState caches, so get_caches_mut returns an empty slice and enforce_max_kv_size_for never trims them (latent inconsistency only, not remotely triggerable today).

…tron_nas/qwen-vl under --max-kv-size (#421) (#422) Completes the offset->live_len prefill-mask migration started in #419/#420 for the remaining dense Vec<KVCache> scheduler-pool models that build their multi-token prefill mask from the cache's monotonic offset. Under the server --max-kv-size path, enforce_max_kv_size_for calls KVCache::trim_front, which advances live_start while offset keeps growing for RoPE. update_and_fetch then returns only live_len = offset - live_start keys, so a mask sized from offset is wider than the live K/V, MLX broadcast_shapes throws, and the non-Result C++ exception crosses the cxx boundary to std::terminate and aborts the whole server. A remote client can crash all sessions with one long prompt plus a multi-token continuation chunk. Fix is mask-only: build the prefill causal mask from cache.live_len() (alias of seq_len(), both = offset - live_start) instead of offset. RoPE, position_ids, and the Llama-4 attention scale keep reading the monotonic offset. Byte-identical when untrimmed (live_start == 0, so live_len == offset). Models changed: - mistral4: transformer_body mask uses caches[0].live_len(); the Llama-4 attention scale below keeps caches[0].offset. - nemotron_nas: forward no-mask branch uses caches.first().map(|c| c.live_len()); per-layer RoPE offset untouched. - qwen2_vl / qwen3_vl / qwen3_vl_moe: only the auto_mask create_causal_mask argument switches to caches[0].live_len(); cache_offset stays monotonic for position_ids/RoPE grid slicing. nemotron_h and jamba were investigated and excluded: they are ModelOwnedSequenceState caches, so get_caches_mut returns an empty slice and enforce_max_kv_size_for never trims them (latent inconsistency only, not remotely triggerable today).

…-max-kv-size (#431) * fix(models): size gemma3/gemma4/exaone_moe prefill mask from live_len gemma3, gemma4, and exaone_moe built their multi-token (l > 1) prefill attention masks from the cache's monotonic offset instead of the live window, for both the global (full-attention) causal mask and the sliding-window prefill mask. That is correct only while no trim has happened. Under the server --max-kv-size path, enforce_max_kv_size_for calls KVCache::trim_front, which advances live_start while offset keeps growing for RoPE; update_and_fetch then returns only live_len = offset - live_start keys, so a mask sized from offset is wider than the live K/V, MLX broadcast_shapes throws, and the non-Result cxx boundary turns the throw into std::terminate and aborts the whole server. This completes the offset->live_len prefill-mask migration started in #419/#420 and #421/#422 for the three Gemma-family / ExaOne dense-cache models they missed. Fix is mask-only: size the prefill masks from the cache's live window (live_len() for a Standard KVCache, seq_len() for a Rotating cache, both equal to the keys update_and_fetch returns) instead of offset. RoPE and position bookkeeping keep reading the monotonic offset. Byte-identical when untrimmed (live_start == 0, so live_len == offset). Models changed: - gemma3: the seq_len > 1 prefill branch sizes the global mask from caches[global_idx].live_len() and the sliding mask from caches[0].live_len(); a live_len() accessor was added to the model's CacheInterface trait (KVCache -> live_len(), RotatingKVCache -> seq_len()). - gemma4: a first_cache_live_len helper (free function plus method form) mirrors first_cache_offset and feeds the two genuine prefill-mask sites (the non-ragged l > 1 branch and the stage-forward execute_hidden path). The batched-MTP sites that use the offset as a per_row_valid_end column coordinate (divergent verify, has_padding, mask_stale_key_gap) intentionally keep the monotonic first_cache_offset; that regime enforces live_len == offset (slot_base == 0) anyway. - exaone_moe: the global and sliding prefill masks size from a new AnyKVCache::live_len() (Standard -> live_len(), Rotating -> seq_len()) instead of offset(). gemma4_mtp_target was verified and needs no change: it is not a registered LanguageModel, drives forward through caller-owned caches that never enter the scheduler pool, and its only mask builder sizes from a fixed offset 0 prefill. Adds per-model unit tests (gemma3_mask_tests, gemma4_unified_mask_tests, exaone_moe_mask_tests) that trim a Standard cache so live_start > 0, then assert the live_len-sized prefill mask key axis matches the continuation chunk's returned K/V (live_len + m) while the offset-sized mask is strictly wider, plus a Rotating-cache case where live_len equals the returned key axis. The shared cache invariant is already pinned by kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim in cache.rs (#419). * fix(models): size gemma3 stage prefill mask from live_len The Gemma3StageModel::execute_hidden pipeline-stage path still sized both the global causal mask and the sliding-window prefill mask from the cache's monotonic offset, while the gemma4 stage analog (Gemma4StageModel::execute_hidden) and the gemma3 single-process path (Gemma3Model::forward_with_caches) were already switched to live_len in this PR. This completes the offset to live_len prefill-mask migration for gemma3 so both stage paths match. Like the rest of #430, this is byte-identical on the untrimmed path: the stage executor's PointerOwnedCacheStore syncs only offset (never live_start) into the model-owned internal caches, so live_start stays 0 and live_len == offset there today. The change uses the existing as_interface().live_len() accessor (KVCache to live_len, RotatingKVCache to seq_len), and keeps RoPE on the monotonic offset. * docs(changelog): record gemma3/gemma4/exaone_moe prefill-mask consistency fix (#430)

The initial v0.3.3 cut (4d5f4fb) was never tagged, and 16 more commits landed on main afterward. Fold those fixes into the v0.3.3 CHANGELOG and debian/changelog entries and set the release date to 2026-06-25: - N-gram loop detection (#433) and Nemotron-H Nano Omni audio input (#443) - Prefill masks sized from the live window under --max-kv-size trim (#418, #420, #422, #431) - Double-transpose crash on mlx-community conv checkpoints (#429) - conv1d/conv2d and nemotron audio-encoder convs fallible at the FFI boundary (#434, #439) - Gemma 4 audio placed in the user turn (#438, #440) - mistral4 MLA backbone routing and 2D MoE token flattening (#423, #425) - Bump actions/checkout from 6 to 7 (#395)

inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Jun 24, 2026

inureyes added status:done Completed and removed status:review Under review labels Jun 24, 2026

inureyes merged commit d3786fd into main Jun 24, 2026
5 checks passed

inureyes deleted the fix/issue-419-dense-mask-live-len-under-trim branch June 24, 2026 10:01

This was referenced Jun 24, 2026

fix: gemma3/gemma4/exaone_moe size prefill mask from monotonic offset, remote-crash under --max-kv-size (missed by #419-#422) #430

Closed

fix: size gemma3/gemma4/exaone_moe prefill mask from live_len under --max-kv-size #431

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(models): size dense prefill masks from live_len under trim (#419)#420

fix(models): size dense prefill masks from live_len under trim (#419)#420
inureyes merged 1 commit into
mainfrom
fix/issue-419-dense-mask-live-len-under-trim

inureyes commented Jun 24, 2026

Uh oh!

inureyes commented Jun 24, 2026

Uh oh!

inureyes commented Jun 24, 2026

Uh oh!

inureyes commented Jun 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

inureyes commented Jun 24, 2026

Summary

The bug

What changed

Why it is correct (and byte-identical when untrimmed)

Test plan

Uh oh!

inureyes commented Jun 24, 2026

Real-model validation across all three dense sliding-window models (before/after)

Uh oh!

inureyes commented Jun 24, 2026

Security / performance review

In scope (the 3 models): correct and complete, no findings

Out of scope, pre-existing (HIGH, recommend follow-up)

Uh oh!

inureyes commented Jun 24, 2026

PR Finalization Complete

Summary

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant