Skip to content

refactor: share the full per-layer attention core across all emitter graph kinds #494

Description

@inureyes

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

  • The shared attention core is used by all three graph kinds: single decode, ragged decode, and prefill.
  • Llama-3.2-1B prefill and ragged assets are byte-for-byte identical before and after (asset test passes).
  • Qwen2 stays token-exact (single-seq) and reference-exact (serve).
  • Gemma2 stays token-exact (single-seq) and reference-exact (serve).
  • Integrated: Llama, Qwen2, and Gemma2 still load and generate via the CLI MLXCEL_BACKEND=xla path and serve via mlxcel-server.

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions