Phase 4 Prefill Optimization

Purpose

Phase 1-3 produced a functionally correct NPU pipeline (kernel-by-kernel in Phase 1, layer-by-layer in Phase 2-3, all numerically aligned with the HF bf16 reference). Phase 4 keeps that correctness while applying the shared optimization skillset to reduce prefill latency. This phase is a thin orchestrator: it dispatches three now-independent optimization skills — opt-merge-multi-launch-kernels (ELF-merging), opt-buffer-object-reuse (host↔NPU runtime-overhead reduction), and opt-layout-alignment (host-side layout) — each of which owns its own recipe and failure modes. These wins are things you compose from the kernels, not behaviors inherited from a reference; the reference's builders are worked examples. Every optimization is an experiment: apply, re-measure, re-run Phase 3 gate; revert if correctness regresses.

For scale, the reference deployment llama3.2-1B took prefill from 18.67 s → 1.30 s (14×) by composing these optimizations — an illustrative datapoint for what they buy on one model, not a target every deployment must hit.

Phase 4 PASS criteria (HARD GATES)

1. Correctness preserved: after every applied optimization, make verify (the token-set gate vs HF bf16) still PASSES — this is the Phase 3 correctness gate, re-run between optimization skills. make diagnosis per-layer cosine is NOT a gate (the verify subsystem retired threshold-based diagnosis; compare_pair reports cosine with no pass/fail); run it only to localize a regression when verify breaks (which optimization / which layer the cosine cliff appears at). If make verify regresses to FAIL, revert the change and document why it doesn't apply to this model.
2. Prefill kernel time strictly < Phase 3 baseline, measured at the same canonical prompt + seq_len with 5-warmup + 20-iter profile. Recording wall time too (kernel + host overhead) is good but the gate is on kernel time (host overhead optimization is a separate concern).
3. Per-skill outcome documented in <model>/docs/development_progress/phase4_prefill.md: for each optimization skill it invoked, record applied / skipped / reverted, the latency delta, and a one-line reason.

The "≥ N optimization skills applied" check is NOT a gate — some models legitimately need only 1-2 (e.g., the model is already seq-first by construction → opt-layout-alignment N/A). The gate is the outcome (perf improved + correctness preserved), not the process count.

Knowledge base references

PRIMARY:

• programming_examples/llms/llama32_1b/docs/profile.md — the reference deployment's profiling breakdown; the reference for what "good" looks like
• programming_examples/llms/<model>/docs/development_progress/phase3_full.md — Phase 3 baseline timings + cosine numbers (the "before" state to measure against and preserve)
• programming_examples/llms/llama_kernel_builder/cache.py — the KernelCache host-optimization knobs (static_input_indices, intermediate_indices); the opt-buffer-object-reuse skill owns how to wire them, this is just the source file it touches

REFERENCE EXEMPLARS (read/mirror to compose your own fused ELFs; import directly only on a bit-for-bit kernel-sequence match):

• programming_examples/llms/llama32_1b/multi_launch_builder/ — the full set of fused-ELF builders, the worked example of how registry leaf kernels stitch into multi-launch ELFs. Mirror these for your model's kernel sequence. Two representative ones:
  • rms_gemms_rope_multi.py — fused 6-launch ELF for RMSNorm + Q/K/V GEMM + RoPE Q/K
  • o_ffn_multi.py — fused 8-launch ELF for O + add + RMSNorm + Gate/Up + SwiGLU + Down + add
• programming_examples/llms/llama_kernel_builder/ — the shared toolkit (KernelCache, stitching, external_kernels) every fused-ELF build uses, inheritance or kernel-first alike.

Workflow

Step 1: Measure Phase 3 baseline

Before invoking any optimization skill, capture the baseline prefill time — this is the number every skill must beat:

cd programming_examples/llms/<model>
flock -x -w 1800 /tmp/mlir-air-npu.lock make profile

Record: kernel time (ms) + wall time as the baseline. Also note the Phase 3 per-layer cosine table as a baseline to compare against if a later optimization breaks verify (it's the localization reference, not a gate).