Repository navigation
refactor: share per-layer attention core across emitter graph kinds - #554
Merged
Merged
Conversation
The OpenXLA emitter already shared the per-layer MLP (`seq_mlp`), but the attention block (q/k/v projection, RoPE, KV read/write, scores, mask, softmax, context, o_proj) was inlined separately in the single decode, ragged decode, and prefill graphs, so a dense architecture that changes attention had to be written once per graph kind. This is issue #494, the Window A foundation that #497 / #498 / #499 / #500 build on. A single driver, `emit_attention`, now emits the complete attention block, driven by an `AttnLayout` enum (`Single` / `Ragged` / `Prefill`) that supplies only the graph-kind-specific layout: activation rank, RoPE broadcast, KV cache indexing, `dot_general` dims, and the softmax axis. The architecture surface (the norm offset, the projection bias, the attention scale, the Gemma2 soft-cap and post-attn norm, and a reserved per-head q/k-norm hook) lives in the driver and its shared free helpers (`arch_norm`, `apply_scale_and_softcap`, `attn_softmax`, `post_attn_norm`), so a new dense family customizes attention once and reaches all three paths together. The uniform-B batched decode is a superseded Stage-1 graph off the serve path and keeps its own inline attention (out of scope, unchanged). The refactor is strictly behavior-preserving: each layout method emits the exact op sequence its path previously inlined, so the emitted StableHLO is byte-for-byte identical. The committed Llama-3.2-1B asset test passes unchanged, and a before/after golden diff over Llama, Qwen2 (tied and untied), and Gemma2 across every graph kind (prefill argmax/logits, decode argmax/logits, ragged b4/b8, batched b4) is byte-identical. A new `gemma2_deltas_reach_every_shared_core_kind` test locks the shared-core guarantee for Gemma2, the one in-scope architecture with no committed byte asset. The public emitter API is unchanged, so the `MLXCEL_BACKEND=xla` CLI and `mlxcel-server` integration compile identical graphs.
inureyes
added
status:done
5 of 6 tasks
inureyes
added a commit
that referenced
this pull request
Jul 1, 2026
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.
5 of 14 tasks
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
…ablup#554) The OpenXLA emitter already shared the per-layer MLP (`seq_mlp`), but the attention block (q/k/v projection, RoPE, KV read/write, scores, mask, softmax, context, o_proj) was inlined separately in the single decode, ragged decode, and prefill graphs, so a dense architecture that changes attention had to be written once per graph kind. This is issue lablup#494, the Window A foundation that lablup#497 / lablup#498 / lablup#499 / lablup#500 build on. A single driver, `emit_attention`, now emits the complete attention block, driven by an `AttnLayout` enum (`Single` / `Ragged` / `Prefill`) that supplies only the graph-kind-specific layout: activation rank, RoPE broadcast, KV cache indexing, `dot_general` dims, and the softmax axis. The architecture surface (the norm offset, the projection bias, the attention scale, the Gemma2 soft-cap and post-attn norm, and a reserved per-head q/k-norm hook) lives in the driver and its shared free helpers (`arch_norm`, `apply_scale_and_softcap`, `attn_softmax`, `post_attn_norm`), so a new dense family customizes attention once and reaches all three paths together. The uniform-B batched decode is a superseded Stage-1 graph off the serve path and keeps its own inline attention (out of scope, unchanged). The refactor is strictly behavior-preserving: each layout method emits the exact op sequence its path previously inlined, so the emitted StableHLO is byte-for-byte identical. The committed Llama-3.2-1B asset test passes unchanged, and a before/after golden diff over Llama, Qwen2 (tied and untied), and Gemma2 across every graph kind (prefill argmax/logits, decode argmax/logits, ragged b4/b8, batched b4) is byte-identical. A new `gemma2_deltas_reach_every_shared_core_kind` test locks the shared-core guarantee for Gemma2, the one in-scope architecture with no committed byte asset. The public emitter API is unchanged, so the `MLXCEL_BACKEND=xla` CLI and `mlxcel-server` integration compile identical graphs.
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
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
Factor the full per-layer attention block into one shared driver reused by the single decode, ragged decode, and prefill emitter graphs, so a future dense family customizes attention once and reaches all three paths together instead of editing each graph kind. Strictly behavior-preserving: the emitted StableHLO is byte-for-byte identical for Llama, Qwen2, and Gemma2 across every graph kind. Window A foundation for #493 (unblocks #497 / #498 / #499 / #500).
What changed
emitter/model.rs: newAttnLayoutenum (Single/Ragged/Prefill) plus anemit_attentiondriver and shared free helpers (arch_norm,apply_scale_and_softcap,attn_softmax,post_attn_norm). The driver sequences the architecture surface (input norm + Gemma2(1+w)offset, q/k/v projection + bias, RoPE, scale, attention soft-cap, mask, softmax, o_proj, Gemma2 post-attn norm, residual) and a reserved per-head q/k-norm hook;AttnLayoutsupplies only the per-kind layout (activation rank, RoPE broadcast, KV cache write/read indexing,dot_generaldims, softmax axis).emitter/model.rs:emit_decode,emit_decode_ragged, andemit_prefillnow build their layout once in the graph head and callemit_attentionin the layer loop, replacing the three inlined attention blocks. The existingseq_mlpsharing and all graph-kind entry points are untouched.emitter/mod.rs: newgemma2_deltas_reach_every_shared_core_kindregression test (with agemma2_like()fixture) locking the shared-core guarantee for Gemma2, the one in-scope architecture with no committed byte asset.Verification
Behavior preservation is proven by byte identity of the emitted graphs, so the compiled vmfbs and
MLXCEL_BACKEND=xlaCLI /mlxcel-serverruntime are identical to before and the prior token-exact / reference-exact results carry over. During development a before/after golden diff over Llama, Qwen2 (tied + untied), and Gemma2 across all seven graph kinds (prefill argmax/logits, decode argmax/logits, ragged b4/b8, batched b4) was byte-for-byte identical; that scratchpad harness is not committed, but the committed asset + structural tests cover the same surface.Test plan
cargo test -p mlxcel-xla --lib(44 passed, incl.from_json_reproduces_bundled_assets_byte_for_byteLlama byte-exact and the new Gemma2 shared-core guard)cargo clippy -p mlxcel-xla --lib --tests -- -D warnings(clean)cargo fmt -p mlxcel-xla(clean)emit_decode/emit_decode_ragged/emit_prefillsignatures and re-export identical), soiree.rsCLI / serve integration compiles identical graphsCloses #494