d5a4ee22ad
The first cut routed every non-shared layer through one causal block mask and relied on SDPA with is_causal=True for shared layers. That gives the right semantics for Gemma-4's full_attention layers but silently drops the sliding_attention window, so sliding layers attend far beyond their 512-token window as soon as the prefix grows. This commit builds a block mask per attention regime (full_attention is pure causal; sliding_attention is causal AND q_pos - kv_pos < window) and passes both into each attention call as a dict, letting the patched forward select by self.layer_type. For the shared-KV sidecar path, the SDPA call now receives an explicit attn_mask composed the same way, so sliding shared layers also respect the window. Strict less-than comparison matches Unsloth's flex-attention convention for GPT-OSS. `_causal_blockmask_with_window` and `_prefill_blockmask_with_window` are local to this file; the shared helpers in `flex_paged_attention.py` stay untouched. `FlexGemma4Inference` now caches both logical decode masks, slices both per-row in `_decode_block_mask`, and runs the PageTable's logical->physical conversion on each before passing them down. |
||
|---|---|---|
| .. | ||
| benchmarks | ||
| enforce_kwargs_spacing.py | ||
| install_gemma4_mlx.sh | ||
| install_qwen3_6_mlx.sh | ||
| run_ruff_format.py | ||