lablup / lablup/mlxcel

fix(cache): bound the turbo4 sliding append to visible_len

Open
#1,793 0 comments 0 reactions 0 assignees View on GitHub
area:core area:inference priority:medium status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem

With `--kv-cache-mode fp16+turbo4` (`KVCacheMode::Turbo4Asym`), the multi-token append of Gemma 4's sliding caches (`update_turbo4_concat`) concatenates the whole physical key buffer, not the `visible_len()` prior keys the FP16 `update_concat` takes. PR #1752 and PR #1776 sized sliding masks from `visible_len()` to match FP16; the turbo append was not changed. After a decode step grows the buffer, or a speculative `trim` rewinds `offset` and `idx` without moving data, it returns `physical + l` keys against an `offset + l` mask and attends zero-filled growth slots and rolled-back draft keys.

Nothing aborts: `trim_mask_to_keys` discards the too-short mask with a warning (one per sliding layer per append) and `causal_attention` rebuilds the band. Whether the extra keys change the attention output is unmeasured. It is reachable with Gemma 4 MTP on a turbo cache: `enable_speculative_buffer` refuses non-FP16 caches, Gemma 4 only logs that, and the next verify block after a rollback hits this. Found during review of PR #1776.

## Evidence

- `src/lib/mlxcel-core/src/cache.rs`: `update_turbo4_concat` `:5003` (physical length `:5028-5031`, whole-buffer concat `:5033-5039`); FP16 `update_concat` `:4686` via `visible_fp16_prefix_for_concat` `:4752`; turbo growth `:5187-5202`; `trim` `:5454-5466`; FP16-only refusal in `enable_speculative_buffer` `:4425-4431`.
- `src/models/gemma4.rs`: per-layer mode in `make_caches_with_modes` `:3801-3820`; refusal only warned at `:5176` and `src/models/gemma4_mtp_target.rs:1717`; `trim_mask_to_keys` discard branch `:2560-2589`.
- The `first_cache_live_len` doc comment PR #1776 added (`gemma4.rs:3923-3946`) defines the live length as the prior keys the next append keeps, true only for FP16, and calls the turbo mask "cropped". It is shorter than the keys, so it is discarded.

## Proposed fix

Take the `visible_len()` prior keys exactly as FP16 does, applying the same slices to `v_packed`, `v_norms` and `v_rescale` (generalize `visible_fp16_prefix_for_concat` rather than copying it). "Exactly" includes chronological order after a wrap (`logical_start`) and keeping `max_size - 1` prior keys plus all new ones (#678), where turbo now clamps the return to `max_size`. Correct the doc comment.

## Acceptance criteria

- [ ] Turbo variants of `first_cache_live_len_sliding_matches_returned_keys` (`gemma4.rs:6095`) and `live_len_matches_the_keys_a_multi_token_append_returns` (`src/models/gemma3_tests.rs:698`), plus a `trim` rollback case, assert that the live length equals the returned key axis, and fail with the fix reverted.
- [ ] No mask/key mismatch warning on a Gemma 4 MTP run with `fp16+turbo4`.
- [ ] A teacher-forced logit trace with and without the fix on `gemma-4-26b-a4b-it-4bit` records whether decided positions change, per `docs/benchmarks.md` "Judging a change that moves the numbers".

## Verification

`examples/logit_trace` cannot reach this as it stands: it builds FP16 caches via `model.make_caches()` (`examples/logit_trace.rs:241`) and runs only multi-token forwards (`:253`, `:257`). The trace arm needs `fp16+turbo4` caches and a single-token step before the traced append, and the reverted arm must log the warning to prove the input reached the branch. Then the workspace gate (`cargo test --workspace --profile test-fast --features metal,accelerate`, clippy, fmt).

Contributor guide

Open the contributing guide

Research direction

Start in src/lib/mlxcel-core/src/cache.rs by comparing update_turbo4_concat with update_concat, visible_fp16_prefix_for_concat, and trim. Run the named Gemma 4 and Gemma 3 cache tests, adding the turbo and rollback cases described in the acceptance criteria; done means returned keys match the live length without mask warnings. Verify the logit trace path and finish with the workspace test, clippy, and fmt checks.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.