Skip to content

feat(xla): add prefill embeddings entry (#858)#879

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-858-prefill-embeddings
Jul 22, 2026
Merged

feat(xla): add prefill embeddings entry (#858)#879
inureyes merged 1 commit into
mainfrom
feature/issue-858-prefill-embeddings

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Adds a distinct prefill_embeddings.main StableHLO entry that accepts post-scale hidden states and an explicit additive attention bias while preserving the existing token-prefill output and KV-cache contract.

Changes

  • Refactors token and embeddings prefill through one shared transformer, normalization, LM-head, output-selection, and KV path.
  • Defines a stable embeddings, positions, real_len, attention_bias schema with fail-closed metadata and finite-mask validation.
  • Preserves tied and untied LM heads, freezes an untied embeddings-prefill golden, and keeps all existing token-prefill goldens byte-identical.
  • Adds a real IREE llvm-cpu parity harness covering tied and untied heads, nonzero padding, real_len = 1, and near-capacity input.

Validation

  • cargo test -p mlxcel-xla --lib prefill_embeddings -- --nocapture (3 passed, 1 opt-in dump test ignored)
  • cargo test -p mlxcel-xla --lib validation::tests::registered_fixtures_are_byte_exact -- --nocapture
  • scripts/xla/validate_arch.sh --structural-only (9 architecture checks passed)
  • cargo check -p mlxcel-xla --lib --tests
  • cargo clippy -p mlxcel-xla --lib --tests -- -D warnings
  • cargo fmt --all -- --check
  • spike/openxla/.venv/bin/python spike/openxla/prefill_embeddings_check.py (exact zero difference for logits and full K/V outputs in every scenario)

Review

Correctness review confirms deterministic argument order, explicit mask polarity and padding semantics, tied/untied head selection, scalar argmax and logits variants, and byte-exact preservation of the token entry. Security review confirms malformed shape, dtype, length, NaN, infinity, and ambiguous mask values fail closed. Performance review found no new copies on the embeddings hot path; validation performs one bounded scan of the static bias before compilation or invocation.

Closes #858

@inureyes inureyes added area:architecture Architecture and code structure changes area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:high High priority status:review Under review type:enhancement New features, capabilities, or significant additions status:done Completed and removed status:review Under review labels Jul 22, 2026
Emit a distinct StableHLO prefill path for post-scale embeddings with an explicit additive attention bias, stable input validation, frozen structural coverage, and real IREE parity checks for tied and untied checkpoints. Refs #858
@inureyes
inureyes force-pushed the feature/issue-858-prefill-embeddings branch from b871e68 to 6c3a23e Compare July 22, 2026 10:07
@inureyes
inureyes merged commit 4b204b3 into main Jul 22, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:architecture Architecture and code structure changes area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:high High priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(xla): add a prefill-from-embeddings StableHLO entry

1 participant