Smart prefill: disaggregating prefill across the Mac fleet

[Gajesh2007]

Summary

Stacks on top of #71 (the embeddings/rerank disaggregated-compute layer). Adds attention-based prompt compression so low-RAM Macs can absorb a meaningful chunk of the prefill load that would otherwise pin the big-RAM fleet.

A small-tier provider runs a tiny draft LLM (default Qwen3-0.6B, 700 MB), captures attention scores over the consumer's prompt, and returns the top-K% of tokens in order. The big-tier provider then runs normal prefill on a 4× shorter prompt with no engine modifications required. Expected: 2-3× lower TTFT today, 5×+ at 128 K context once we graduate to true sparse prefill via vllm-mlx#179 in phase 2.

This is the answer to "how do we disaggregate prefill across consumer internet?" The naive answer (ship KV cache) is dead — 8 GB per 32 K-token Qwen 27B prefill vs 125 MB/s residential pipe. The right answer is don't ship KV, ship a shorter prompt. Same compute leverage, ~6 orders of magnitude less bandwidth.

Linked issue

Closes #

What ships

Two surfaces, one shared dispatch path:

1. POST /v1/compress — explicit. Useful for pre-compressing a RAG corpus once and storing the compressed chunks.
2. smart_prefill field on /v1/chat/completions — opt-in middleware. Coordinator dispatches the longest user/system message to a tiny-tier compressor, swaps the result back in, then routes to the consumer's chosen big model. Surfaces stats via X-SmartPrefill-* response headers.

Both call into the same runCompression helper so retry, billing, and E2E encryption are guaranteed identical.

Wire protocol (coordinator/internal/protocol/messages.go ↔ provider/src/protocol.rs)

• New message pair: prompt_compression_request / prompt_compression_complete
• E2E encrypted under the same NaCl box session-key flow as embeddings/chat — coordinator never sees plaintext on either leg
• Round-trip tests on both sides

Tier routing (coordinator/internal/registry/registry.go)

• New compressor model_type joins embedding / rerank in PreferredTiersForModelType — routes to tiny/small tier first, falls back to standard

Catalog (coordinator/cmd/coordinator/main.go + provider/src/main.rs)

• mlx-community/Qwen3-0.6B (700 MB, 8 GB min RAM)
• mlx-community/Qwen3-1.7B (1.8 GB, 8 GB min RAM, higher quality)
• Cross-family drafts work (arXiv 2603.02631) → a single 0.6 B serves the entire big-tier catalog. No per-target compressors needed.

Pricing (coordinator/internal/payments/pricing.go)

• $0.004 / 1 M tokens (Qwen3-0.6B), $0.006 / 1 M (1.7B)
• Net economic impact for the consumer: ~19× positive — compression cost is dwarfed by the prefill cycles it saves on the big-tier model
• Same 95 / 5 provider / platform split

Provider proxy (provider/src/proxy.rs, coordinator.rs, main.rs)

• New handle_compression_request forwards to a local HTTP sidecar at EIGENINFERENCE_COMPRESSOR_PORT (default embedding_port + 1) exposing POST /v1/compress
• Same E2E session-key flow as embeddings — provider seals the compressed prompt back over the coordinator's session pubkey
• A first-party MLX-based sidecar will ship with the next provider bundle; until then the proxy will return 502 (connect refused) and the coordinator will route to another tiny provider via the existing retry loop

Docs (docs/smart-prefill.md) — full design, citing the recent literature: SpecPrefill (NVIDIA/MS 2025), Cross-Family Speculative Prefill (arXiv 2603.02631, March 2026), PrfaaS (arXiv 2604.15039, April 2026), BEAVER (arXiv 2603.19635, March 2026), and the open vllm-mlx#179 prototype on our exact runtime that demonstrates 5.45× TTFT reduction at 128 K.

Review-round-2 fixes

Two independent reviewers caught real bugs in the first revision; commit 6de0a851 fixes all of them:

Bug	Severity	Fix	Regression test
smart_prefill field leaked into chat-provider request body when middleware fell through AND consumer set explicit max_tokens	Real (vllm-mlx would reject the unknown field on the fall-through path)	Always re-marshal rawBody after applySmartPrefill — the middleware always strips the field, so the re-marshal must be unconditional too	TestSmartPrefillStripsFieldOnFallThrough
Reservation not refunded when compression succeeded but produced empty result or swap failed → consumer billed for compression they got no benefit from	Real	Single refundCompression closure called from every post-billing fall-through path (compression error, empty result, swap failure)	TestSmartPrefillRefundsOnEmptyResult
max_keep_tokens defined in protocol but never plumbed through middleware settings	NIT	Added to smartPrefillSettings, parsed from object form, forwarded to compressor request body	covered by structural tests

Reviewer 2 also questioned the cleanup pattern in dispatchCompression's decrypt-failure path; verified against the embeddings/rerank canonical flow — it's correct (pending was already removed by handlePromptCompressionComplete before the channel send, decrypt failure only needs SetProviderIdle + RecordJobFailure, both of which are called). Reviewer 1 independently reached the same conclusion.

Test plan

• cd coordinator && go test ./... — all green (8 new tests, including 2 regression tests)
• cd coordinator && gofmt -l . — clean
• cd coordinator && go vet ./... — clean
• rustfmt --check --edition 2024 on the Rust files I touched — clean

New tests:

• TestCompressE2E — full encrypted round-trip on the standalone endpoint through a 16 GB tiny provider
• TestSmartPrefillMiddlewareSwapsLongestMessage — full encrypted round-trip through the middleware: register a tiny compressor + a standard chat provider, assert the chat provider receives the compressed prompt and the response carries the X-SmartPrefill-* headers
• TestSmartPrefillFallsThroughOnShortPrompt — middleware no-ops below the threshold; chat provider sees the original prompt
• TestSmartPrefillStripsFieldOnFallThrough — regression test for review-round-2 fix DGInf: Decentralized private AI inference platform #1
• TestSmartPrefillRefundsOnEmptyResult — regression test for review-round-2 fix feat: build Mac private inference core #2, with the billing service active
• TestCompressNoFreeCreditWhenBillingDisabled — same regression test as embeddings (refund cannot mint balance when billing was never charged)
• TestCompressInvalidRatio — input validation
• TestPreferredTiersIncludesCompressor — tier preference wired up
• 4 new Rust protocol round-trip tests in provider/src/protocol.rs

Components touched

• coordinator (Go)
• provider (Rust)
• console-ui (Next.js)
• image-bridge (Python)
• app (macOS Swift)
• enclave (Swift)
• infra / CI / release
• docs

Protocol / interface changes

• Yes — described above and matching side updated

Both sides updated symmetrically; round-trip tests on both sides catch drift. The smart_prefill field on /v1/chat/completions is a new request extension; we strip it from the body before forwarding to the provider so the existing OpenAI-compatible providers never see it.

No bundle script changes — the compressor sidecar runs at a deterministic port (embedding_port + 1 or EIGENINFERENCE_COMPRESSOR_PORT) so the launcher and the proxy agree without extra IPC.

Notes for reviewers

Why now and why this design:

I asked Exa for the latest literature before writing a line of code. The field moved past LLMLingua-2 — the new winner is attention-based "Speculative Prefill" using a small draft LLM (NVIDIA/MS 2025), with cross-family drafts proven viable in March 2026. There's an open prototype of exactly this on the vllm-mlx fork our providers run (waybarrios/vllm-mlx#179) that reports 5.45× TTFT reduction on Qwen3.5-122B at 128 K context on M2 Ultra. So this PR ships the protocol envelope + tier routing + billing now, and the next iteration graduates from text-level compression (phase 1, this PR) to true sparse prefill against the original positional schema (phase 2, follow-up). Clients don't change between phases — the same smart_prefill flag governs both.

What this is NOT:

• Not bit-exact prefill disaggregation. That requires shipping KV cache, which is bandwidth-bound across consumer internet and dead on arrival.
• Not speculative decoding. That accelerates decode using a draft on the same machine. Both will be supported, neither replaces the other.
• Not lossless. Phase 1 drops ~75 % of input tokens. Quality on LongBench/RULER at 4× is consistently >90 % across the cited research, but consumers needing verbatim recall (legal, code, exact-quote retrieval) should leave smart_prefill off.

Stacking note: This PR's base branch is cursor/disaggregated-compute-small-providers-40ff (PR #71), not master. The diff against #71 is what this PR adds; the diff against master is the union. Once #71 lands, GitHub will auto-rebase this against master.