Repository navigation
fix: reset model-owned VLM fallback state - #34
Merged
Merged
Conversation
VLM single-row generation paths (CLI and `bench` invocations that do not carry a `SequenceId`) were leaking fallback cache state between runs. `CxxGenerator::generate_streaming{,_with_embeddings}` and `generate_with_stats{,_and_embeddings}` only called `self.reset()`, which clears the generator-owned `KVCache` vector but never touches the model-owned fallback slot inside `ModelOwnedSequenceState`. A second VLM prefill then reused offsets from the previous run against a freshly built 4D attention mask, producing wrong outputs or panicking when the cached position fell outside the new mask.
Introduce a new `LanguageModel::reset_runtime_state()` trait hook with a default no-op body (correct for models whose entire state lives in the caller's `KVCache` slice) and call it from `CxxGenerator::reset_with_model`. Switch every single-row generation path off the bare `reset()` and onto `reset_with_model(model)` so the generator-owned caches and the model-owned fallback caches are cleared in lockstep. Per-sequence scheduler entries are intentionally untouched — only the legacy fallback slot used by `SequenceId`-less paths is reset.
Per-model implementations: Gemma 3, Gemma 4, and Llama 4 forward the hook to their existing `reset_caches()` routines. Qwen 3.5 Next instead calls `sequence_state.replace_internal(make_internal_caches())` directly so the VLM-populated MRoPE fallback slot is preserved — VLM callers write that slot during embedding preparation immediately before generation begins. `Gemma4VLModel` and `Qwen35VLModel` delegate the hook to their text backbones, and both `VisionLanguageModel` and `LoadedModel` propagate it through their existing `delegate_language_model!` machinery so every wrapper layer participates.
Add `tests/vlm_concurrency.rs` — an ignored, weight-gated smoke test that drives two concurrent image chat requests against a real `mlxcel-server` Gemma 3n VLM instance and asserts both responses come back without cross-contamination. Add a `reset_with_model_invokes_runtime_state_hook` unit test in `mlxcel-core` that uses a tracking model to assert `reset_with_model` actually invokes the new hook, so future single-row paths cannot silently regress past it.
inureyes
added
type:bug
3 tasks
inureyes
added a commit
that referenced
this pull request
May 20, 2026
`qwen2.5-vl-3b-4bit` was marked FAIL:warmup in the M5 Max benchmark sweep. The failure surfaced in the same-process bench harness: a warmup pass left stale generator-owned KV state, and the subsequent measured pass broadcast a 4-D VLM attention mask against those stale keys and aborted. PR #34 (`fix: reset model-owned VLM fallback state`) fixed this transitively by rebuilding generator-owned KV caches before each single-row run in `CxxGenerator::reset_with_model()`; `qwen2_5_vl` carries no model-owned fallback state, so no additional production code change is needed here. This commit adds the regression test `qwen2_5_vl_single_row_bench_prefill_succeeds_with_m5_warmup` in `tests/vlm_concurrency.rs`, which drives the exact FAIL:warmup repro path — an image VLM run with `--warmup-tokens 20` so the harness performs both a warmup pass and a measured pass in the same process. It mirrors the existing `gemma3n_vlm_single_row_bench_prefill_accepts_m5_tile_padding` guard and is `#[ignore]`d behind local Qwen2.5-VL weights and the `mlxcel-bench-decode` binary. The M5 Max benchmark doc (`docs/benchmark_results/model_tests_m5max.md`) is updated to reflect the now-passing run: text prefill 1529.43 tok/s / decode 165.11 tok/s (EOS at 39 tokens) and VLM prefill 1696.53 tok/s / decode 156.83 tok/s (EOS at 22 tokens, 1.61x vs M1 Ultra). The pass count rises from 91 to 92, the fail count drops from 5 to 4, qwen2.5-vl-3b-4bit leaves both the Known Issues table and the mlx-lm unsupported-architecture list, and it gains a real mlx-lm decode comparison row (156.83 vs 98.53 tok/s, 159%).
7 tasks
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
VLM single-row generation paths (CLI and `bench` invocations that do not carry a `SequenceId`) were leaking fallback cache state between runs. `CxxGenerator::generate_streaming{,_with_embeddings}` and `generate_with_stats{,_and_embeddings}` only called `self.reset()`, which clears the generator-owned `KVCache` vector but never touches the model-owned fallback slot inside `ModelOwnedSequenceState`. A second VLM prefill then reused offsets from the previous run against a freshly built 4D attention mask, producing wrong outputs or panicking when the cached position fell outside the new mask.
Introduce a new `LanguageModel::reset_runtime_state()` trait hook with a default no-op body (correct for models whose entire state lives in the caller's `KVCache` slice) and call it from `CxxGenerator::reset_with_model`. Switch every single-row generation path off the bare `reset()` and onto `reset_with_model(model)` so the generator-owned caches and the model-owned fallback caches are cleared in lockstep. Per-sequence scheduler entries are intentionally untouched — only the legacy fallback slot used by `SequenceId`-less paths is reset.
Per-model implementations: Gemma 3, Gemma 4, and Llama 4 forward the hook to their existing `reset_caches()` routines. Qwen 3.5 Next instead calls `sequence_state.replace_internal(make_internal_caches())` directly so the VLM-populated MRoPE fallback slot is preserved — VLM callers write that slot during embedding preparation immediately before generation begins. `Gemma4VLModel` and `Qwen35VLModel` delegate the hook to their text backbones, and both `VisionLanguageModel` and `LoadedModel` propagate it through their existing `delegate_language_model!` machinery so every wrapper layer participates.
Add `tests/vlm_concurrency.rs` — an ignored, weight-gated smoke test that drives two concurrent image chat requests against a real `mlxcel-server` Gemma 3n VLM instance and asserts both responses come back without cross-contamination. Add a `reset_with_model_invokes_runtime_state_hook` unit test in `mlxcel-core` that uses a tracking model to assert `reset_with_model` actually invokes the new hook, so future single-row paths cannot silently regress past it.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
Bump version to 0.0.29. Highlights: model-owned VLM fallback state reset (lablup#34), Gemma 3n VLM padded prefill alignment plus the new mlxcel-bench-decode same-process benchmark harness (lablup#36), and benchmark report reorganization with refreshed Apple Silicon and M5 Max VLM tables (lablup#32, lablup#33).
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
`qwen2.5-vl-3b-4bit` was marked FAIL:warmup in the M5 Max benchmark sweep. The failure surfaced in the same-process bench harness: a warmup pass left stale generator-owned KV state, and the subsequent measured pass broadcast a 4-D VLM attention mask against those stale keys and aborted. PR lablup#34 (`fix: reset model-owned VLM fallback state`) fixed this transitively by rebuilding generator-owned KV caches before each single-row run in `CxxGenerator::reset_with_model()`; `qwen2_5_vl` carries no model-owned fallback state, so no additional production code change is needed here. This commit adds the regression test `qwen2_5_vl_single_row_bench_prefill_succeeds_with_m5_warmup` in `tests/vlm_concurrency.rs`, which drives the exact FAIL:warmup repro path — an image VLM run with `--warmup-tokens 20` so the harness performs both a warmup pass and a measured pass in the same process. It mirrors the existing `gemma3n_vlm_single_row_bench_prefill_accepts_m5_tile_padding` guard and is `#[ignore]`d behind local Qwen2.5-VL weights and the `mlxcel-bench-decode` binary. The M5 Max benchmark doc (`docs/benchmark_results/model_tests_m5max.md`) is updated to reflect the now-passing run: text prefill 1529.43 tok/s / decode 165.11 tok/s (EOS at 39 tokens) and VLM prefill 1696.53 tok/s / decode 156.83 tok/s (EOS at 22 tokens, 1.61x vs M1 Ultra). The pass count rises from 91 to 92, the fail count drops from 5 to 4, qwen2.5-vl-3b-4bit leaves both the Known Issues table and the mlx-lm unsupported-architecture list, and it gains a real mlx-lm decode comparison row (156.83 vs 98.53 tok/s, 159%).
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.
Summary
VLM single-row generation paths (CLI and
benchinvocations that do not carry aSequenceId) were leaking fallback cache state between runs.CxxGenerator::generate_streaming{,_with_embeddings}andgenerate_with_stats{,_and_embeddings}only calledself.reset(), which clears the generator-ownedKVCachevector but never touches the model-owned fallback slot insideModelOwnedSequenceState. A second VLM prefill then reused offsets from the previous run against a freshly built 4D attention mask, producing wrong outputs or panicking when the cached position fell outside the new mask.Fix
Introduce a new
LanguageModel::reset_runtime_state()trait hook with a default no-op body (correct for models whose entire state lives in the caller'sKVCacheslice) and call it fromCxxGenerator::reset_with_model. Switch every single-row generation path off the barereset()and ontoreset_with_model(model)so the generator-owned caches and the model-owned fallback caches are cleared in lockstep. Per-sequence scheduler entries are intentionally untouched — only the legacy fallback slot used bySequenceId-less paths is reset.Per-Model Implementations
reset_caches()routines.sequence_state.replace_internal(make_internal_caches())directly so the VLM-populated MRoPE fallback slot is preserved; VLM callers write that slot during embedding preparation immediately before generation begins.delegate_language_model!machinery so every wrapper layer participates.Test plan
cargo fmt --checkcargo check --tests --features metal,acceleratecargo clippy --features metal,accelerate --all-targets -- -D warningsreset_with_model_invokes_runtime_state_hookunit test (uses a tracking model to guaranteereset_with_modelactually invokes the new hook, so future single-row paths cannot silently regress past it)tests/vlm_concurrency.rs— ignored, weight-gated smoke test driving two concurrent image chat requests against a realmlxcel-serverGemma 3n VLM instance; run locally withcargo test --release --test vlm_concurrency -- --ignoredwhen Gemma 3n VLM weights are available