Repository navigation
test(xla): validate a LLaVA reference architecture end to end #862
Copy link
Copy link
Closed
Labels
area:benchmarkBenchmark harness and performance measurement (bench_*.sh, /update-benchmarks)Benchmark harness and performance measurement (bench_*.sh, /update-benchmarks)area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)Generation, sampling, decoding (incl. speculative, DRY)priority:highHigh priorityHigh prioritystatus:doneCompletedCompletedtype:testTest related changesTest related changes
Description
Parent and dependency
- Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
- Blocked by feat(xla): admit image requests in CLI and continuous-batch serving #861
- Existing XLA validation entry points:
scripts/xla/validate_arch.sh,xla_oracle_check, andxla_batch_bench - Design reference:
spike/openxla/MULTIMODAL_VLM_DESIGN.md
Problem
The first multimodal path crosses processor, token/placeholder expansion, text embedding, vision projection, attention bias, StableHLO, IREE, KV seeding, decode, detokenization, and server streaming. Unit tests at any one seam cannot prove the assembled path matches the reference model. The epic needs a reproducible architecture-level gate before additional VLM families copy the pattern.
Goal
Qualify one pinned LLaVA-style reference checkpoint end to end against an independent reference implementation, through both CLI and OpenAI-compatible server paths.
Reference fixture contract
- Select and document one publicly obtainable, license-compatible checkpoint/revision whose text architecture is already supported by the XLA emitter and whose CLIP/SigLIP plus projector path is supported by refactor(vlm): extract a host multimodal preprocessor for XLA #859.
- Pin repository, revision, tokenizer/processor files, expected image size, dtype, and deterministic test image/prompt. Do not rely on a moving default branch.
- Keep large weights and generated artifacts out of Git; the harness must cache/download them through existing repository conventions and clearly report an unavailable network/device prerequisite.
Required implementation
- Add a checked-in oracle harness that obtains processor outputs, expanded tokens/positions/mask, projected image features, first prefill logits, selected KV slices, and a deterministic greedy token sequence from a maintained independent implementation such as Hugging Face Transformers.
- Add a matching mlxcel XLA runner that captures the same checkpoints. Compare intermediate values before comparing final text so a failure identifies processor, vision/projector, LM prefill, KV, or decode drift.
- Define per-stage absolute/relative tolerances by dtype. Token output must be exact for deterministic greedy decoding; tolerance-only final validation is insufficient.
- Exercise at least: image plus text, two images if the reference architecture supports it, no-image text fallback, malformed placeholder count, and context-capacity overflow.
- Run one non-streaming CLI request and one streaming OpenAI-compatible server request. Validate finish reason, chunk ordering, detokenized content, and usage/metrics semantics.
- Add the reference architecture to
scripts/xla/validate_arch.shor a clearly linked companion command and document required IREE target/device settings. - Record compile time, peak host preprocessing memory, runtime device memory, prefill latency, and decode throughput as non-regression baselines. Performance does not need to beat MLX in this issue, but pathological duplicate-model memory is a failure.
Validation matrix
- Pure structural/emitter tests that run without IREE.
- CPU IREE execution when supported in CI or the documented local gate.
- At least one production-relevant Metal or CUDA target, with the exact target recorded in the PR.
- Existing text-only architecture oracle and batch validation to prove no regression.
Acceptance criteria
- The checkpoint, revision, prompt, image, processor settings, and tolerances are pinned and reproducible.
- Processor, projected features, prefill logits, selected all-layer KV data, and greedy output are compared to an independent oracle.
- Greedy output is token-exact for the documented fixture.
- Both CLI and streaming server requests pass, including usage and finish semantics.
- Failure output identifies the first divergent stage.
- A real IREE accelerator target passes in addition to structural tests.
- Memory measurements prove the host-first path does not retain a duplicate full text decoder.
- The validation command is documented as the merge gate for later VLM-family issues.
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
area:benchmarkBenchmark harness and performance measurement (bench_*.sh, /update-benchmarks)Benchmark harness and performance measurement (bench_*.sh, /update-benchmarks)area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)Generation, sampling, decoding (incl. speculative, DRY)priority:highHigh priorityHigh prioritystatus:doneCompletedCompletedtype:testTest related changesTest related changes
Activity
test(xla): validate a LLaVA reference architecture end to end (#897)