You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Gemma2 alternates sliding-window attention (window 4096) and global attention across layers. The emitter currently implements global attention only, which is a no-op difference for prompts shorter than the window but produces wrong results once the context exceeds 4096 tokens.
Scope
Implement the per-layer sliding-window mask so the local/global alternation matches Gemma2.
Wire the window size and the per-layer local/global selection from Config::from_json.
Files: emitter/model.rs, emitter/config.rs.
Acceptance criteria
Token-exact vs the HF fp32 oracle on a prompt longer than the window. Validated at reduced scale: the emitter caps context at MAX_SEQ = 256 (< the 4096 production window), so the check runs the identical mask code with a shrunk window that bites inside 256. A synthetic Gemma2 (matched HF config) is token-exact and bit-exact (max |logit diff| ~6e-8) vs the HF fp32 oracle at both an inert window (>= prompt) and a biting window (< prompt), and HF's own output changes between the two, so the biting case is not vacuous. The literal >4096-token GPU sweep needs MAX_SEQ > 4096 (out of scope here) and is the orchestrator's post-merge gate.
Existing short-prompt Gemma2 outputs unchanged (token-exact single-seq, reference-exact serve). For W >= MAX_SEQ the window predicate is always true, so the local mask equals the causal mask on the value; the inert-window oracle run is bit-exact vs HF, and the Llama-3.2-1B byte-for-byte asset test is unchanged.
Integrated: the sliding-window mask is emitted in all three graph kinds the CLI (MLXCEL_BACKEND=xla) and mlxcel-server compile from config.json at load (single decode, prefill, ragged decode), so a Gemma2 checkpoint runs with the local/global alternation on both paths. Long-context behaviour beyond 256 tokens is bounded by MAX_SEQ (a separate follow-up), not by the masking.
Validation
Pure-Rust structural tests (scoped cargo test -p mlxcel-xla --lib emitter::tests): window parse + default, is_sliding_layer even-local schedule, one-time local-mask block per graph, window value emitted, and per-layer alternation (even layers consume the local mask, odd the global, and they differ). Llama byte-for-byte asset test still passes (non-Gemma2 emission byte-identical).
Part of #493
Context
Gemma2 alternates sliding-window attention (window 4096) and global attention across layers. The emitter currently implements global attention only, which is a no-op difference for prompts shorter than the window but produces wrong results once the context exceeds 4096 tokens.
Scope
Config::from_json.emitter/model.rs,emitter/config.rs.Acceptance criteria
MAX_SEQ = 256(< the 4096 production window), so the check runs the identical mask code with a shrunk window that bites inside 256. A synthetic Gemma2 (matched HF config) is token-exact and bit-exact (max |logit diff| ~6e-8) vs the HF fp32 oracle at both an inert window (>= prompt) and a biting window (< prompt), and HF's own output changes between the two, so the biting case is not vacuous. The literal >4096-token GPU sweep needsMAX_SEQ > 4096(out of scope here) and is the orchestrator's post-merge gate.W >= MAX_SEQthe window predicate is always true, so the local mask equals the causal mask on the value; the inert-window oracle run is bit-exact vs HF, and the Llama-3.2-1B byte-for-byte asset test is unchanged.MLXCEL_BACKEND=xla) andmlxcel-servercompile fromconfig.jsonat load (single decode, prefill, ragged decode), so a Gemma2 checkpoint runs with the local/global alternation on both paths. Long-context behaviour beyond 256 tokens is bounded byMAX_SEQ(a separate follow-up), not by the masking.Validation
cargo test -p mlxcel-xla --lib emitter::tests): window parse + default,is_sliding_layereven-local schedule, one-time local-mask block per graph, window value emitted, and per-layer alternation (even layers consume the local mask, odd the global, and they differ). Llama byte-for-byte asset test still passes (non-Gemma2 emission byte-identical).spike/openxla/gemma2_sliding_window_check.py(emit -> IREE llvm-cpu -> compare to HF eager fp32).Dependencies
None required. Benefits from #494 (shared attention core) if it lands first, but can proceed independently.