Tile the privacy filter's local attention band instead of a dense [L, L] mask - #76
Open
eauchs wants to merge 1 commit into
Open
Tile the privacy filter's local attention band instead of a dense [L, L] mask#76eauchs wants to merge 1 commit into
eauchs wants to merge 1 commit into
Conversation
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.
openai_privacy_filterbuilds its local attention band as a dense[L, L]mask, shared across layers. That is O(L²): at 32k tokens the mask alone is 2 GB in bf16, plus the int32 index grid used to build it.Above
attention_tile_threshold(default 4096) the band is now resolved one query tile at a time. Each tile carries every key its queries can reach, soscaled_dot_product_attentionis still called with the fused kernel and each softmax, sink included, is exactly the one the dense path computes. Below the threshold nothing changes.M3 Max 128 GB,
openai/privacy-filterbf16, single sequence, minimum of 3 runs:Throughput becomes flat in sequence length instead of collapsing: 6.6x at 32k, with peak memory down 2.8x. The model advertises a 128k context, which the dense mask cannot reach.
Correctness: logits are bit-identical to the dense path (
max |diff| = 0) at 300 / 1,500 / 6,000 tokens for tile sizes 256, 512 and 1,024, and on padded batches for every tile size down toL/4. Two tests added alongside the existingtest_openai_privacy_filter_model, covering tiled-vs-dense equality and padding.attention_tile_sizedefaults to 512, the fastest of 512 / 1,024 / 2,048 at every length measured. The threshold is deliberately conservative: 512-token tiles already beat the dense mask at 2,048 tokens, so it could be lowered, but batches of short sequences keep the existing path untouched this way.Pre-existing unrelated failure on
test_qwen3_vl_model_process_uses_high_level_processor_paths, present onmainas well; the other 51 tests pass.