Skip to content

feat: implement Gemma2 sliding-window attention (local/global alternation) #495

Description

@inureyes

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

  • 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).
  • Independent HF-oracle execution check: 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.

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions