Repository navigation
Conversation
A RotatingKVCache returns window - 1 + L keys for an L-token chunk, and the array mask built for it made the fused kernel visit every one of them for each query row. Process query bands of 256 rows over only the keys their window can reach, starting on a 64-key boundary so the kernel's tile grid and online softmax order are unchanged. The result is bit-identical to the single call. Chunks of 8 rows or fewer (and tails of 8 rows or fewer) keep the single call, since the fused kernel hands those to sdpa_vector, which accumulates in a different order. Head dims other than 64 and 128 run the unfused fallback, whose softmax reduction depends on the key count, so they keep the single call too (gemma-3 uses 256). Caches with keep > 0 and masks that are not [L, keys] are left on the old path.
The banded path inferred a sliding-window mask from a [L, kL] array mask on a keep=0 RotatingKVCache and took its window from cache.max_size. A mask built with a wider window_size, or any other [L, kL] mask, then lost keys it allows: max abs error 0.38 for window 4096 on a 513-slot cache and 0.36 for an all-True mask. scaled_dot_product_attention now takes an optional sliding_window, and bands only when it equals cache.max_size, i.e. when the caller states the mask is that cache's own window mask. Flint passes it for its sliding layers; every other model keeps the single SDPA call.
The shared scaled_dot_product_attention cannot tell whether a mask is a causal sliding-window mask, or which window it uses, so banding there is only correct when every caller passes the right window. Do it in Flint's attention instead, with the band bounds taken from the model's own sliding_window, and leave base.py unchanged.
tc3oliver
changed the title
tc3oliver
changed the title
This branch has not been deployed
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
Flint’s sliding-window prefill processes fully masked key tiles in SDPA.
Split prefill into 256-query bands, skipping unreachable keys. Key starts are aligned to 64 to preserve the fused kernel’s tile boundaries.
Only Flint’s sliding-attention layers are affected. Unsupported cases fall back to the original SDPA. Shared base.py is unchanged.
Benchmarks
M4 Max (40-core GPU, 64 GB), MLX 0.32.4, RecursiveSI/flint-int4. 2048-token chunks, median of three alternating runs.
Prefill throughput, not server TTFT. Last-token logits MD5 matched.
Correctness
Outputs match the original SDPA exactly (mx.array_equal) across tested cases:
Head dimension 256, short tails, and unsupported configurations use the original path.
Tests
python -m unittest discover tests:
Both errors are from missing optional dependencies (datasets, lm_eval). Black, isort, Ruff, and mutation checks passed.
Tested on M4 Max; M5 not tested.
Reproduction
Script: ab2.py