Moonshot: fused single-query embedding megakernel for MiniLM

[Davidobot]

Moonshot / research spike. Filed to capture the idea and its gate, not as committed work.

Motivation

For a single query, embedding dominates end-to-end latency. On M1 (docs/benchmarks.md, all-MiniLM-L6-v2, 384-dim): full query p50 is 5.73 ms (CPU embed) to 8.42 ms (MPS embed), while the TurboVec scan is ~0.25 ms p50. The scan is already sub-millisecond and fuses scoring with top-k; the embedding forward pass is the remaining lever for single-query p50/p95.

A batch-1 transformer-encoder forward pass is the regime megakernel fusion targets: many small dependent ops where per-launch overhead and HBM round-trips of activations between ops dominate, not FLOPs. Published results in this regime: Hazy Research's Llama-1B megakernel sustains ~78% of H100 memory bandwidth at bs=1; Mirage MPK reports 1.2-6.7x latency reductions by compiling the whole forward pass into one kernel.

Idea

Fuse the entire all-MiniLM-L6-v2 forward pass into a single GPU kernel for batch=1: token and position embed, then [LayerNorm, QKV, attention, FFN] x 6, mean-pool, L2-normalize, keeping activations on-chip across all six layers with no intermediate global-memory writes and no per-op launch. Query-only fast path (CUDA and/or Metal); batched and bulk embedding stay on the ONNX/torch runtime.

Why this is a moonshot, not planned work

• Off-mission. LodeDB is a vector store; the embedding runtime is opt-in and deliberately delegated to ONNX Runtime (feat(embeddings): default to ONNX Runtime with a PyTorch fallback #29, feat(embeddings): make the embedding runtime an opt-in extra #48). A hand-fused encoder kernel is inference-engine surface we chose not to own.
• Maintenance surface. One kernel per (architecture, hidden size, layer count, precision, backend). This buys MiniLM only; BGE, CLIP, and every other preset would each need their own.
• Correctness bar. Output must match the default runtime within tolerance to preserve recall, including pooling (mean vs CLS) and normalization. The Swift NLEmbedder identity-guard work (Swift/iOS binding over the native core: feature parity + docs (#28) #47) shows how easily embedding drift slips in.
• Portability. CUDA and Metal are separate implementations, and the default single-query laptop path is CPU/NEON, which gains nothing from a GPU kernel.

Spike gate (before writing any kernel)

1. Attribute the 5.7-8.4 ms: how much is the model forward pass vs tokenization, ORT session/dispatch, and host/device copies. If the forward pass is not the majority, the ceiling is low.
2. Check whether an existing compiler already captures most of the win with zero custom-kernel maintenance: ONNX Runtime graph fusion / TensorRT, or a Mirage-MPK-style auto-megakernel. A megakernel we hand-own is justified only if it beats the best compiler path by a multiplicative margin.
3. Prototype one fused encoder layer, measure, extrapolate to six.

Proceed to a full megakernel only if the spike shows a durable multiplicative win over the tuned ONNX path.

Constraints

• Query-only opt-in fast path; ONNX default and bulk embedding unchanged.
• Embedding output within parity tolerance of the default runtime (recall-preserving).
• No change to storage formats.

Related: #27 (resident core / persistent-kernel serving), #29, #48 (embedding-runtime opt-in).