inference_efficiency

Inference-efficiency experiments (Co-LMLM vs LMLM-Asker)

Measures the per-retrieval query-formation overhead of the two retrieval mechanisms: LMLM-Asker decodes a natural-language question and encodes it with a separate sentence-transformer, whereas Co-LMLM reads its query vector from the decoder hidden state (a single <FACT> forward). The FAISS search itself is excluded (shared by both, index-placement dependent).

All model/data paths are passed via config/args — nothing is baked into these scripts. Everything runs via uv run. Copy the example configs and fill in your own paths.

Canonical method — per-fact overhead on the dynamic-ppl gold set

Free-form generation confounds the comparison: the two models trigger different numbers of retrievals, at different positions, with different content and per-lookup context lengths. Instead we measure at fixed gold <FACT> positions of the dynamic-perplexity eval set, so each model forms a query from the same context. For each fact we warm the KV cache on the context prefix (prefill excluded), then time only the query-formation step.

• dynppl_overhead.py — the harness (pyrallis config, see configs/dynppl.example.yaml). Per fact, with a warm cache (vLLM prefix caching):
  • retriever: w1 = latency to forward <FACT> (Co-LMLM query vector)
  • asker: wq = generate the question until </QUESTION>, we = jina encode (bf16+FA2)
• run_dynppl.sh <gpu> <config.yaml> [overrides…] — thin wrapper around the pyrallis script.
• analyze_dynppl.py <out_dir> [--decode-ms M --encode-fwd-ms M] — reads {retriever,asker}.jsonl from <out_dir> and reports the overhead pure-compute (fixed per-call invocation overhead removed for both) and as-measured.

Run (copy configs/dynppl.example.yaml → configs/dynppl.local.yaml, fill in paths):

bash run_dynppl.sh 0 configs/dynppl.local.yaml \
     --phase retriever --model_path <retriever_ckpt> --output_path out/dynppl/retriever.jsonl
bash run_dynppl.sh 1 configs/dynppl.local.yaml \
     --phase asker     --model_path <asker_ckpt>     --output_path out/dynppl/asker.jsonl
uv run --no-sync python analyze_dynppl.py out/dynppl

Result (FW-360M, medians over ~1.5K facts, pure-compute): Co-LMLM forms its query in a single ~2.2 ms <FACT> forward; LMLM-Asker spends ~27 ms decoding the ~13-token question plus ~1 ms for the encoder forward (~28 ms total) — a ~13× reduction in per-retrieval overhead. The asker's cost is dominated by decoding the query: the encoder forward itself is cheap (the raw single-query .encode() wall is ~46 ms but that is ~all fixed invocation overhead, see calibration below).

Calibration

The pure-compute numbers use forward/decode compute, subtracting fixed per-call overhead; these benchmarks establish the rates (analyze_dynppl.py defaults match FW-360M bf16):

• decode_rate.py <model> — vLLM steady-state decode (ms/token) via a time = a + b·N linear fit (FW-360M bf16 ≈ 2.185 ms/tok; per-call overhead ≈ 6.5 ms).
• hf_decode_rate.py <model> [attn_impl] — HF eager-loop decode rate cross-check (confirms the ~10× HF/vLLM gap is the un-graphed loop, not the attention kernel; ~22 ms/tok).
• jina_encode_rate.py [--model id] — batch-scales the jina encoder to split the ~46 ms single-query .encode() wall into fixed invocation overhead (~45 ms) vs. per-query forward compute (~1.1 ms, the saturated-regime marginal; ≤32 queries encode in the same wall as one).

Alternative — generation-time per-phase decomposition

run_timed.sh <vllm|hf> <asker|retriever> <gpu> runs the actual retrieval generation with the per-phase timers in the generators (t_generate_s / t_encode_s / t_search_s / gen_decoded_tokens, emitted per record); analyze_timed.py decomposes it. This measures overhead during real generation (confounded by differing retrieval counts/positions, which is why the dynppl method above is canonical), but exercises the instrumentation and includes the FAISS phase. Paths come from env vars — copy config.example.sh → config.local.sh, fill it in, then source config.local.sh && bash run_timed.sh vllm asker 0.