Skip to content

Tile the privacy filter's local attention band instead of a dense [L, L] mask - #76

Open
eauchs wants to merge 1 commit into
Blaizzy:mainfrom
eauchs:banded-attention-tiling
Open

Tile the privacy filter's local attention band instead of a dense [L, L] mask#76
eauchs wants to merge 1 commit into
Blaizzy:mainfrom
eauchs:banded-attention-tiling

Conversation

@eauchs

@eauchs eauchs commented Sep 3, 2026

Copy link
Copy Markdown

openai_privacy_filter builds 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, so scaled_dot_product_attention is 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-filter bf16, single sequence, minimum of 3 runs:

tokens main this PR peak memory
2,048 43,396 tok/s 50,055 tok/s 3.09 → 3.02 GB
8,192 27,003 tok/s 55,252 tok/s 3.41 → 3.22 GB
32,768 8,012 tok/s 53,030 tok/s 10.62 → 3.82 GB

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 to L/4. Two tests added alongside the existing test_openai_privacy_filter_model, covering tiled-vs-dense equality and padding.

attention_tile_size defaults 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 on main as well; the other 51 tests pass.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant