Part of #493
Context
The OpenXLA emitter already shares the per-layer MLP (seq_mlp), but the attention block (q/k/v projection, RoPE, KV read/write, scores, mask, softmax, context, o_proj) is still inlined separately in each graph kind: single decode, ragged decode, and prefill. A new dense architecture that changes the attention block (for example adding per-head q/k normalization) must therefore be written two or three times, once per graph kind. This is the main reason adding a dense family is not yet trivial.
Scope
- Factor the attention core into a shared, shape-parametric form, or a per-graph-kind callback driven by a shared layer driver, so a dense architecture's attention is authored once and reaches the single-sequence and serve paths together.
- Keep the existing MLP sharing (
seq_mlp) and the graph-kind entry points intact.
- Files:
emitter/model.rs (primary), with supporting changes in emitter/config.rs / iree.rs kept minimal.
Acceptance criteria
Implementation notes
A single driver, emit_attention, now emits the complete per-layer 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, softmax axis). The architecture surface (norm offset, projection bias, attention scale, Gemma2 soft-cap and post-attn norm, and a reserved per-head q/k-norm hook) lives in the driver and its shared helpers, 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).
Behavior preservation is proven by byte identity of the emitted StableHLO: the committed Llama-3.2-1B asset test passes unchanged, and a before/after golden diff over Llama, Qwen2 (tied + untied), and Gemma2 across every graph kind (prefill argmax/logits, decode argmax/logits, ragged b4/b8, batched b4) is byte-for-byte identical. Because the emitted graphs are identical and the public emitter API is unchanged, the compiled vmfbs and the MLXCEL_BACKEND=xla CLI / mlxcel-server runtime behavior are identical to before, so the prior token-exact / reference-exact results carry over.
Dependencies
None. This is the Window A foundation that unblocks Windows B and C.
Part of #493
Context
The OpenXLA emitter already shares the per-layer MLP (
seq_mlp), but the attention block (q/k/v projection, RoPE, KV read/write, scores, mask, softmax, context, o_proj) is still inlined separately in each graph kind: single decode, ragged decode, and prefill. A new dense architecture that changes the attention block (for example adding per-head q/k normalization) must therefore be written two or three times, once per graph kind. This is the main reason adding a dense family is not yet trivial.Scope
seq_mlp) and the graph-kind entry points intact.emitter/model.rs(primary), with supporting changes inemitter/config.rs/iree.rskept minimal.Acceptance criteria
MLXCEL_BACKEND=xlapath and serve viamlxcel-server.Implementation notes
A single driver,
emit_attention, now emits the complete per-layer attention block, driven by anAttnLayoutenum (Single/Ragged/Prefill) that supplies only the graph-kind-specific layout (activation rank, RoPE broadcast, KV cache indexing,dot_generaldims, softmax axis). The architecture surface (norm offset, projection bias, attention scale, Gemma2 soft-cap and post-attn norm, and a reserved per-head q/k-norm hook) lives in the driver and its shared helpers, 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).Behavior preservation is proven by byte identity of the emitted StableHLO: the committed Llama-3.2-1B asset test passes unchanged, and a before/after golden diff over Llama, Qwen2 (tied + untied), and Gemma2 across every graph kind (prefill argmax/logits, decode argmax/logits, ragged b4/b8, batched b4) is byte-for-byte identical. Because the emitted graphs are identical and the public emitter API is unchanged, the compiled vmfbs and the
MLXCEL_BACKEND=xlaCLI /mlxcel-serverruntime behavior are identical to before, so the prior token-exact / reference-exact results carry over.Dependencies
None. This is the Window A foundation that unblocks Windows B and C.