Skip to content

feat(xla): support Phi4MM audio and per-slot speech adapters #877

Description

@inureyes

Parent and dependencies

Problem

Phi4MM audio requires both a cascaded Conformer encoder/projection and modality-specific decoder LoRA. A global adapter switch is unsafe for continuous batching: text, speech, vision, and mixed requests can occupy different slots simultaneously, while one XLA module bundle shares immutable weights.

Moving only the encoder to IREE while omitting speech LoRA would fail the qualified #874 architecture. Mutating one process-wide adapter between requests would create cross-request correctness and isolation bugs.

Goal

Execute the qualified Phi4MM audio encoder, projection, placeholder replacement, speech/vision adapter selection, language prefill, and decode through XLA/IREE for CLI and continuous-batch serving.

Required implementation

  • Emit an audio.main module for the pinned cascaded Conformer and projection contract from feat(vlm): add a token-exact Phi4MM audio reference path #874. Accept host-extracted feature frames plus explicit valid-frame masks/lengths; use static frame buckets and reject/bucket before compilation/invocation.
  • Port every encoder primitive that affects inference, including subsampling/length propagation, attention/relative positions, convolution, normalization, feed-forward, downsample grouping, and speech-versus-vision projection branch.
  • Produce projected audio tokens and valid output lengths in the exact order expected by the canonical prepared-prefill placeholder mapping. Keep concatenation/replacement cardinality checks on both Rust and native boundaries.
  • Implement modality LoRA in the XLA language graph without global mutable state. Store one adapter/input mode per batch slot and pass it through prefill and every decode step. A ragged/mixed batch must apply base-only, speech, or vision deltas independently per row.
  • Load complete base plus LoRA A/B pairs with rank/shape/scale validation and bind adapter configuration into artifact identity. Missing one half of a required pair is a load error, not a warning.
  • Choose a graph strategy that preserves per-slot correctness (for example batched low-rank delta with row gating or explicitly isolated engine pools). If isolated pools are used, document the loss of cross-mode batching and prevent a request from crossing pools mid-generation.
  • Keep the mode selected during prefill for decode. Cancellation/slot reset must clear mode and audio buffers; reused slots default to no adapter until assigned.
  • Integrate feat(xla): add bounded audio preprocessing and serving plumbing #873 CLI/server audio ingestion and the feat(vlm): add a token-exact Phi4MM audio reference path #874 prompt/multi-clip rules. XLA capability is true only when audio module, required adapter modules/weights, and text bundle are all compatible.
  • Do not retain a full MLX Phi4MM decoder. Host feature extraction is allowed; encoder/projection/LM math must not silently fall back to MLX.

Non-goals

Validation

  • Compare feature input, selected Conformer intermediates, projected audio tokens/lengths, merged embeddings, LoRA-adjusted QKV/O/MLP outputs at selected decoder layers, logits/KV, and greedy tokens against feat(vlm): add a token-exact Phi4MM audio reference path #874 and the official oracle.
  • Interleave text, speech, vision, and mixed requests in one scheduler workload. Prove each slot's adapter remains stable through decode and clears on cancellation/reuse.
  • Add negative tests for global adapter leakage, missing/incomplete LoRA pairs, wrong projection mode, frame-bucket overflow, length/placeholder mismatch, and a request routed to an incompatible engine pool.
  • Run audio-only and mixed-media CLI plus streaming server tests on a production-relevant IREE target; final deterministic ASR output must be token-exact.
  • Run all Phi4MM image/text and general XLA continuous-batch regression gates.

Acceptance criteria

  • Phi4MM Conformer/projection executes through IREE and matches feat(vlm): add a token-exact Phi4MM audio reference path #874 intermediates.
  • Adapter/input mode is per slot, persists through decode, and cannot leak across mixed batches.
  • Complete LoRA pairs/scales are validated and included in artifact compatibility.
  • Projected-length and placeholder cardinality are checked at every boundary.
  • Audio and mixed-media deterministic outputs are token-exact through CLI/server.
  • No full decoder duplication or silent MLX encoder fallback remains.
  • Existing text/image XLA behavior remains token-exact.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:architectureArchitecture and code structure changesarea:inferenceGeneration, sampling, decoding (incl. speculative, DRY)area:modelsModel architectures, weights, loading, metadatapriority:lowLow prioritystatus:blockedBlocked by dependencies or other issuestype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions