Repository navigation
fix: apply KV modes to model-owned caches - #1400
Merged
Merged
Conversation
Model-owned families were building their internal attention caches with default FP16 constructors, so the resolved KV mode only affected dense external cache slices and left Gemma 3/4, Qwen 3.5, Qwen3-Next, LFM2, Bailing MoE Linear, AFMoE, and Llama 4 regular attention arms unchanged. Add model-owned per-layer KV-mode hooks, inject the already-resolved effective layer table from the CLI generator and server scheduler, and teach the affected model-owned constructors and VLM wrappers to consume it while keeping recurrent, convolutional, gated-delta, and Llama 4 chunked state at FP16. Make prompt-cache and snapshot reuse conservative across KV modes by rejecting stale dense detached cache modes with a named reason and refusing non-FP16 model-owned snapshots unless their mode matches the live configured table. Validation: - cargo check -p mlxcel --lib - cargo check -p mlxcel --tests - cargo test --profile test-fast -p mlxcel-core reset_with_model_forwards_resolved_modes_to_model_owned_caches -- --nocapture - cargo test --profile test-fast -p mlxcel --lib qwen3_next::cache_tests -- --nocapture - cargo test --profile test-fast -p mlxcel --lib dense_detached_kv_mode_validation -- --nocapture - cargo build -p mlxcel --bin mlxcel --profile test-fast - target/test-fast/mlxcel generate -m models/gemma3-1b-4bit -p 'Write a detailed 300 word essay about the history of Roman aqueducts, their engineering, and their legacy.' -n 200 -t 0 --kv-cache-mode fp16 - target/test-fast/mlxcel generate -m models/gemma3-1b-4bit -p 'Write a detailed 300 word essay about the history of Roman aqueducts, their engineering, and their legacy.' -n 200 -t 0 --kv-cache-mode int8 Refs #1330
inureyes
added
status:review
The cache stats route added kv_cache_mode_effective for issue #1330, but it still read only the legacy kv_cache_mode field. When server-side --kv-bits enabled BatchKvQuantConfig, the scheduler applied Int8 or Turbo4Asym while /v1/cache/stats continued to report fp16. Prefer the batch KV base mode whenever batch quantization is enabled, and cover the Int8, TurboQuant, and legacy fallback cases with a route helper test. Validation: - cargo test --profile test-fast -p mlxcel --lib cache_stats_effective_kv_mode_prefers_batch_kv_quant_when_enabled -- --nocapture Refs #1330
Member
Author
Security/performance review after 0177304Findings
Validation
Safety audit
|
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
## Summary Apply the resolved effective KV cache mode table to model-owned attention caches so `--kv-cache-mode` and `--kv-bits` no longer stop at dense external cache slices, while keeping observability and prompt-cache reuse aligned with the mode that is actually active. ## What changed - Added `LanguageModel` KV-mode table hooks plus a shared `KvCacheLayerModes` helper and delegated the hooks through loaded-model and VLM wrappers. - Injected resolved per-layer modes from the CLI generator and server scheduler, including Boundary-V and batch KV quantization policy, then reported the applied-layer count in the CLI and server surfaces. - Updated Gemma 3, AFMoE, Gemma 4, Qwen 3.5, Qwen3-Next, Bailing MoE Linear, LFM2, and Llama 4 regular attention caches to use the injected mode while keeping recurrent, convolutional, gated-delta, and Llama 4 chunked caches unchanged. - Made dense prompt-cache adoption reject stale per-layer KV modes with `kv_mode_mismatch`, and made model-owned snapshot restore refuse non-FP16 or mode-mismatched attention state instead of silently truncating quantization sidecars. - Added `kv_cache_mode_effective` to cache statistics, preferring the batch KV quantization base mode when `--kv-bits` is active and otherwise using the legacy effective mode. - Documented model-owned KV-mode behavior, snapshot conservatism, prompt-cache mismatch observability, and the Llama 4 `ChunkedKVCache` exception. ## Changes during review - Fixed `/v1/cache/stats` to report the batch KV quantization mode instead of incorrectly falling back to `fp16` when `BatchKvQuantConfig` is enabled. ## Validation - [x] `cargo test --workspace --profile test-fast --features metal,accelerate` - [x] `cargo clippy --workspace --all-targets -- -D warnings` - [x] `cargo fmt --all -- --check` - [x] Focused generator forwarding, Qwen3-Next snapshot, dense prompt-cache mismatch, and cache-stats effective-mode tests - [x] Gemma 3 real-checkpoint generation on Metal with both `fp16` and `int8` Real checkpoint result: Gemma 3 FP16 and INT8 both completed on Metal, the banners reported `applied to 26 of 26 layers`, stderr was empty, and the greedy outputs differed, confirming that the model-owned cache consumed the configured mode. Closes lablup#1330
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
Apply the resolved effective KV cache mode table to model-owned attention caches so
--kv-cache-modeand--kv-bitsno longer stop at dense external cache slices, while keeping observability and prompt-cache reuse aligned with the mode that is actually active.What changed
LanguageModelKV-mode table hooks plus a sharedKvCacheLayerModeshelper and delegated the hooks through loaded-model and VLM wrappers.kv_mode_mismatch, and made model-owned snapshot restore refuse non-FP16 or mode-mismatched attention state instead of silently truncating quantization sidecars.kv_cache_mode_effectiveto cache statistics, preferring the batch KV quantization base mode when--kv-bitsis active and otherwise using the legacy effective mode.ChunkedKVCacheexception.Changes during review
/v1/cache/statsto report the batch KV quantization mode instead of incorrectly falling back tofp16whenBatchKvQuantConfigis enabled.Validation
cargo test --workspace --profile test-fast --features metal,acceleratecargo clippy --workspace --all-targets -- -D warningscargo fmt --all -- --checkfp16andint8Real checkpoint result: Gemma 3 FP16 and INT8 both completed on Metal, the banners reported
applied to 26 of 26 layers, stderr was empty, and the greedy outputs differed, confirming that the model-owned cache consumed the configured mode.Closes #1330