feat: add Gemma2 sliding-window (local/global) attention#555
Merged
Conversation
Gemma2 alternates sliding-window (local) and full (global) attention across layers, but the OpenXLA emitter previously emitted global attention on every layer. That is correct only while the context stays under the window; once it exceeds the window the local layers must stop attending to older keys, so long-context Gemma2 output was wrong. This adds the per-layer sliding-window mask so the local/global alternation matches Gemma2 (issue #495). Config gains an Option<usize> sliding_window, read from a gemma2 checkpoint's config.json (HF default 4096) and left None for Llama/Qwen2 (whose own sliding_window field is ignored, matching use_sliding_window = false). is_sliding_layer encodes the Gemma2 schedule: even layers local, odd layers global. The three graph kinds the CLI and server compile at load (single decode, prefill, ragged decode) build a local mask once by anding the window into the existing causal mask with one extra select (within-window keeps the causal value, else -1e30), and each local layer broadcasts that mask instead of the global one. The local mask is emitted only for a windowed config, so Llama/Qwen2 graphs are byte-identical (the bundled Llama-3.2-1B asset test still passes), and for a window >= MAX_SEQ it is a value no-op, so short-context Gemma2 output is unchanged. Validated by pure-Rust structural tests (window parse and default, the even-local schedule, one local-mask block per graph, the window constant emitted, and the per-layer alternation) plus an independent HF-oracle execution check (spike/openxla/gemma2_sliding_window_check.py): a matched synthetic Gemma2 is token-exact and bit-exact (max abs logit diff ~6e-8) against HF eager fp32 at both an inert window and a window small enough to bite inside the 256-slot cache, with HF's own output changing between the two.
inureyes
added
type:enhancement
origin/main landed #554 (issue #494), which unified the three attention emit paths (single decode, ragged decode, prefill) into one shared `emit_attention` driver parameterized by an `AttnLayout` enum, applying the additive key mask in a single shared step. This branch (issue #495) had added Gemma2 sliding-window (local/global alternation) to the pre-#494 three separate emit paths, so the two changes collided in `emitter/model.rs`. Resolved by fitting #495's per-layer mask choice into #494's shared core: each `AttnLayout` variant now carries an optional `mask_local` (the sliding-window mask, built once per graph head only when a window is configured), and the shared `add_mask` step selects local vs global per layer through a single `layer_mask`/`is_sliding_layer` check. All three layouts (Single, Ragged, Prefill) get the alternation uniformly from one authoring site. Non-windowed architectures (Llama, Qwen2) parse `sliding_window = None`, so no window ops are emitted and `layer_mask` returns the global mask handle unchanged: the emitted graphs stay byte-for-byte identical to origin/main, confirmed by the Llama byte-exact asset guard across all five graph kinds. Also unified the two `gemma2_like` test helpers (one added by each side) into a single windowed helper that keeps #494's non-square o_proj coverage and #495's four-layer alternation. emitter::tests: 17 passed, 1 ignored; clippy and fmt clean.
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
1 participant
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
Implements Gemma2's alternating sliding-window (local) / full (global) attention in the OpenXLA emitter (issue #495). The emitter previously emitted global attention on every layer, which is correct only while the context stays under the window; the local (even) layers must stop attending to keys older than the window. The mask is wired from
Config::from_json, so a Gemma2 checkpoint runs with the correct local/global alternation on both the single-sequence CLI (MLXCEL_BACKEND=xla) and themlxcel-serverserve path, since both compile the same emitter graphs at load.What changed
emitter/config.rs: newOption<usize> sliding_windowfield, read from a gemma2config.json(HF Gemma2 default 4096), leftNonefor Llama/Qwen2 (Qwen2's ownsliding_windowis ignored, matchinguse_sliding_window = false). Newis_sliding_layer(li)schedule: even layers local, odd global (matches HFGemma2DecoderLayer).emitter/model.rs: in each of the three graph kinds (single decode, prefill, ragged decode), build a local mask once by anding the window into the existing causal mask with one extraselect(within-window keeps the causal value, else-1e30); each local layer broadcasts that mask, each global layer the unchanged causal mask. Emitted only for a windowed config, so Llama/Qwen2 graphs stay byte-identical; a value no-op when the window>= MAX_SEQ, so short-context Gemma2 output is unchanged.emitter/mod.rs: structural tests plus an opt-in (#[ignore]) graph-dump test used by the execution check.spike/openxla/gemma2_sliding_window_check.py: independent HF-oracle execution check (emit -> IREE llvm-cpu -> compare to HF eager fp32).Validation
Numbers are from the local run; the full >4096-token GPU sweep is bounded by
MAX_SEQ = 256(a separate follow-up) and is left for the post-merge gate.cargo test -p mlxcel-xla --lib emitter::tests: 16 passed (1 ignored). Includes the Llama-3.2-1B byte-for-byte asset test (non-Gemma2 emission unchanged), the even-local schedule, the one-time local-mask block per graph, the emitted window constant, and the per-layer alternation (even layers consume the local mask, odd the global, and they differ).spike/openxla/gemma2_sliding_window_check.py: a matched synthetic Gemma2 is token-exact and bit-exact vs HF eager fp32 at an inert window (max |logit diff| 4.5e-8) and at a window small enough to bite inside the 256-slot cache (5.9e-8); HF's own output changes by 4.7e-2 between the two, so the biting case is not vacuous.Test plan
cargo check -p mlxcel-xla --lib --testscargo clippy -p mlxcel-xla --lib --tests -- -D warningscargo test -p mlxcel-xla --lib emitter::testsspike/openxla/gemma2_sliding_window_check.py(HF fp32 oracle, inert + biting window)Closes #495