Add pplx-embed (Perplexity) bidirectional Qwen3 encoder on the ANE

[dokterbob]

Adds a stateless bidirectional Qwen3-0.6B encoder path for Perplexity's pplx-embed embedding models, running 99.8% on the ANE. Two variants: plain sentence embeddings and late chunking (per-chunk via a pool_matrix matmul, one encoder pass over the whole window).

Conversion (conversion/)

• models/qwen3_encoder.py — ANE encoder reusing the existing primitives (Conv2d-1×1 projections, repeat_kv_ane, stable_attention): QK-norm, RoPE θ=1e6, SwiGLU, GQA 16/8, full bidirectional pad-mask. Derives seq-len + batch from the input. fp16 residual rescale (K=8) keeps the 28-layer down_proj accumulation in range (exact for a pre-norm net). RMSNorm is now a local Qwen3RMSNorm selected via norm_impl (see workstream A).
• build_pplx_embed_bundle.py — fixed-shape ANE buckets (the fast path) + an optional flexible RangeDim GPU model (--dynamic-upper N) as the >max-bucket catch-all; plain + context; native int8 / pooled_fp16 output.
• pplx_embed_reference.py — fp32 golden oracle, bit-exact with the model's own st_quantize.py (int8/binary/ubinary). test_pplx_embed_parity.py gates it.
• config.py — registry entries pplx-embed, pplx-embed-context.

Swift (Sources/)

• CoreMLLLM/PplxEmbed.swift — [String] → int8/binary/ubinary (plain) and per-chunk context; tokenize → smallest fitting bucket → ANE, with n > largest bucket routed to the flexible GPU model (non-padded). Adds PplxEmbed.load(repo:buckets:preferCompiled:) for manifest-driven download (see workstream C). pplx-embed-demo CLI (--repo) + pplx-embed-bench harness.

Fidelity (cosine vs fp32 reference)

plain int8 ≥0.999, context ≥0.998, dynamic GPU 0.9996. Swift-vs-Python end-to-end 0.9998.

Design decision: fixed buckets only on the ANE

EnumeratedShapes/RangeDim force CPU fallback on the ANE and are ~10× slower (measured: 969 ms @512 vs 101 ms). So the ANE path is one fixed-shape .mlpackage per bucket (pad to the smallest fitting); the flexible RangeDim model is GPU-only and used solely as the >max-bucket catch-all.

Side investigations (reproducible scripts + write-ups in docs/)

1. Quantization — docs/PPLX_EMBED_W8A8.md + conversion/experiment_w8a8.py. Weight-only int8 collapses (~0.42 cos), int4-palettize only reaches 0.905, full W8A8 collapses to ~0. ANE residency stays fine, so it's purely numerical — and moot anyway since int8 activations are only ~9% faster here (not bandwidth-bound). Conclusion: ship fp16 + buckets; recovery would need full QAT.
2. Batching throughput — docs/PPLX_EMBED_BATCHING.md + conversion/experiment_batching.py. Batching is not a useful throughput lever: the ANE is batch-1 by design (batching hurts, 0.69×, confirmed 100% on-ANE), the GPU path batches only modestly (~1.4× at small L) and saturates by L≈512. Batch-1 on the ANE at the smallest bucket is the throughput winner. (Measured on an Apple M4 Max.)

---

What changed (this session)

Four workstreams on top of the encoder path above. All measurements: Apple M4 Max, macOS 26, coremltools 9, torch 2.11, B=1.

A — Native RMSNorm, shipped as the default

A new local Qwen3RMSNorm (x * rsqrt(mean(x²) + eps) * w) plus a norm_impl config selecting it at the 5 encoder norm sites, leaving the shared ane_ops.ANERMSNorm (the cat([x, −x]) → LayerNorm → chunk trick) untouched. The A/B (conversion/experiment_ane_rmsnorm.py, changing only the norm sites and holding Conv2d-1×1 projections + tensor layout fixed) isolated the RMSNorm as the cause of the ANE speedup seen incidentally during the earlier GPU-residency investigation. The cat/chunk trick predated a fast native ANE rsqrt; on this stack it is now a de-optimization. norm_impl="native" is now the pplx-embed default. Resolves the open follow-up flagged in docs/PPLX_EMBED_GPU_RESIDENCY.md. Shared-decoder rollout (the ~10 other families on ane_ops.ANERMSNorm) is flagged but not done — needs per-family re-validation — in docs/ANE_RMSNORM_FOLLOWUP.md.

B — Larger fixed ANE bucket (L=8192): measured, does NOT ship

conversion/measure_l8192_bucket.py shows a fixed L=8192 bucket statically maps 99.81% of ops to the ANE, but the ANE runtime fails at inference (ANEProgramProcessRequestDirect status=0x15): the 8192² full-attention intermediates (16 heads × 8192² fp16 ≈ 2 GB per score tensor) exceed ANE buffer limits, and the ANE compile alone takes ~25 min. So the largest fixed ANE bucket stays 4096; 4097–8192 tokens keep the existing dynamic GPU catch-all.

C — Single fixed RoPE table → fully shared weights

The encoder now builds its RoPE cos/sin tables at a single fixed size (max_position_embeddings) for every bucket, applied via a runtime position_ids gather (fold-proof — a static [:S] slice gets const-folded back to per-bucket). The only per-bucket-varying constant is gone, so weight.bin is now byte-identical across all buckets. Output is bit-identical to the old per-bucket RoPE (parity + dynamic-RangeDim + Swift int8 all verified). This is what makes the distribution below dedup.

D — Hugging Face distribution with end-to-end dedup

conversion/upload_pplx_embed.py publishes the prebuilt buckets to one HF repo — per-bucket subfolders + manifest.json, shipping both .mlmodelc (default — no on-device compile) and .mlpackage per bucket. Because all buckets now share one weight blob:

• Upload (~2.3 GB, not ~14 GB): conversion/upload_pplx_embed_dedup.py uploads each unique blob once and server-side-copies the duplicate paths (CommitOperationCopy) — sidestepping hf upload-large-folder's re-upload of identical oids across its parallel batches.
• Download (~1.15 GB, not ~3.5 GB for the default 3 buckets): PplxEmbed.load(repo:buckets:preferCompiled:) uses HF's native Swift Hub client (swift-huggingface HubClient.downloadSnapshot, with the Xet trait — HF stores large files Xet-backed by default) into the content-addressed cache, so the shared weight is fetched once no matter how many buckets you pull. Selective per-bucket, one format per bucket, exact-path matching; feeds the existing load(bundleDir:). pplx-embed-demo --repo exercises it end to end.

Verification

• A: native RMSNorm is +12.7% @l256 / +21.5% @L512 on CPU_AND_NE vs ane_cat, at identical 99.81% ANE residency and cosine 0.99998 vs the fp32 oracle. Parity gate PASS.
• B: L=8192 static plan = 99.81% ANE but ANE inference fails (status=0x15); ~25 min compile. Decision: do not ship; 4097–8192 stays on the GPU catch-all.
• C: single fixed RoPE table → weight.bin byte-identical across buckets (verified by sha256); parity PASS, dynamic-RangeDim + Swift int8 bit-identical to the old encoder.
• D: end-to-end download verified live — pplx-embed-demo --repo dokterbob/pplx-embed-coreml --buckets 512 returns correct embeddings, and the content-addressed cache holds one 1.1 GB blob with two weight.bin paths pointing to it (download dedup confirmed). Upload deduped via server-side copy.

Follow-ups

• Shared-decoder native RMSNorm — roll norm_impl="native" to the shared ane_ops.ANERMSNorm only per-family, after re-validating decode + prefill latency, residency, and parity (docs/ANE_RMSNORM_FOLLOWUP.md). Do not flip the shared default globally off the pplx-embed result alone.
• B3 — mMARCO calibration + multilingual retrieval eval (nDCG@10 across languages).
• chunk-and-pool for >8192 tokens — stay within a fixed ANE bucket instead of the GPU catch-all.

Uses the existing pip/venv conversion workflow (a separate PR adds an optional uv manifest). Verified: test_pplx_embed_parity.py PASS, swift build clean.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.