feat: add NIAH eval to cache_rate_tester

[ziqifan617]

Motivation

At high concurrency, KV cache corruption or block-reuse bugs can cause the model to emit wrong or garbage
output — a failure mode that is completely invisible in TTFT and throughput metrics alone. This PR makes it
detectable with zero external dependencies and no separate eval phase.

What this adds

An optional in-band output correctness eval for cache_rate_tester.py, based on the passkey retrieval task
introduced in Landmark Attention: Random-Access Infinite Context Length for Transformers (Mohtashami & Jaggi,
2023).

When --eval-mode niah is enabled, a configurable fraction of working-set prompts (default 10%) are replaced
with needle-in-a-haystack probes: a coherent English haystack of the same context length as the rest of the
test, with a random N-digit passkey embedded at a random position (10–90% depth), followed by a retrieval
question. Eval prompts are interleaved with the regular synthetic prompts throughout the timed run, so they
exercise the same KV cache behavior and the same concurrency pressure as every other request.

Each response is graded by substring match against the expected passkey. With greedy decoding and a healthy
cache, any modern model trivially retrieves the passkey — a sustained drop below 100% accuracy is a strong
signal that something is wrong with the inference path, not noise.

We use 7-digit passkeys by default, giving a false-positive rate of ~10⁻⁷ per response — effectively zero
across hundreds of eval probes per run.

New flags

Flag	Default	Description
--eval-mode {none,niah}	none	Enable NIAH eval. none preserves all existing behavior.
--eval-fraction FLOAT	0.1	Fraction of working-set prompts replaced with eval probes.
--eval-passkey-digits INT	7	Digits in the random passkey. Higher = lower false-positive rate.
--eval-output-tokens INT	512	Output token budget for eval prompts. Thinking models (e.g. Qwen3) consume 200–400 tokens in <think> reasoning before writing the answer, so eval prompts need a larger budget than regular requests without inflating the perf-measurement budget.

Output

detailed_results_*.csv — per-request columns on every eval probe:

• eval_expected — the passkey that should appear in the response
• eval_passed — True / False
• eval_response_excerpt — first 300 chars of the response (for debugging failures)

sustained_periods_*.csv — per 30-second assessment window:

• eval_total, eval_passed, eval_accuracy — shows when during the run accuracy degraded

summary_*.csv — single aggregate across all requests in the run:

• eval_total, eval_passed, eval_accuracy

Eval grading only activates at cache_hit_rate=100. At mixed cache rates the tester appends random
gibberish to each prompt to drive the desired miss fraction, which would clobber the retrieval question; this
is enforced automatically.

Sample run

   Period 3: Running at concurrency 32 for 30.0s
      Prefills: 7, Contributing: 7, Launched: 9
      Input: 25,641 tok/s | Output: 1,510 tok/s (streaming-based)
      Avg TTFT: 12.676s | P95 TTFT: 16.847s | P99 TTFT: 17.203s
      Avg ITL: 22.18ms | avg_output_tokens: 374.1 tok/s
      Eval: 7/7 passed (100.0%)
      Concurrency 32 = max-concurrency → MAX_REACHED

    [... periods 4–18 omitted for brevity, all Eval: X/X passed (100.0%) ...]

    Period 19: Running at concurrency 32 for 30.0s
      Prefills: 2, Contributing: 2, Launched: 3
      Input: 21,463 tok/s | Output: 1,265 tok/s (streaming-based)
      Avg TTFT: 16.441s | P95 TTFT: 19.832s | P99 TTFT: 20.104s
      Avg ITL: 24.61ms | avg_output_tokens: 391.2 tok/s
      Eval: 2/2 passed (100.0%)
      Concurrency 32 = max-concurrency → MAX_REACHED

  ================================================================================
  ✓ All tests complete!
  ================================================================================
  Results saved to: /tmp/cache_rate_qwen3-32b-tp1-kvbm-v2-eval_c32_g1g2

  Total continuous tests completed: 1

  ═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════
  ═════════════════
  Final Summary - All Test Results
  ═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════
  ═════════════════
     Context   Cache%   Requests    Input Tok   Output Tok      Input/s     Output/s   Avg TTFT    Conc   EvalAcc
  ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
  ─────────────────
      32,000     100%         437       12.71M        0.75M      23,309        1,372     13.410s      32    100.0%
  ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
  ─────────────────
       TOTAL              437          12.71M        0.75M
  ═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════
  ═════════════════

[callanjfox]

Thanks for this — really nice addition. The design is clean: in-band probes that share the same load and concurrency as regular requests, deterministic seed-driven selection, correct gating to cache_hit_rate=100, and the CSV/summary plumbing all look good. It's fully additive too — --eval-mode none (the default) preserves every existing path. I'd like to land it. 👍

One blocker — needs a rebase. The branch currently conflicts with master. Since you opened this, master picked up two prompt-path optimizations ("Cache decoded prompt text to avoid re-decoding full context per request" and "Pre-generate unique-token pool…") that rewrote the exact region this PR touches — the cache_hit_rate == 100 branch of construct_prompt where you add the eval_metadata = working_set.get_eval_metadata(session_id) lookup. The resolution is mechanical: re-apply the eval-metadata lookup and the "skip QUESTION_BANK for eval prompts" logic on top of the new decode-caching code. Could you rebase onto latest master and push?

Non-blocking follow-ups (happy to take these as separate issues if you'd rather land first):

1. --eval-output-tokens default is inconsistent. argparse uses default=1024 and TestConfig uses 1024, but the flag's help text says "(default: 512)" and the PR description table also says 512. Actual default is 1024 — let's reconcile (I'd keep 1024 and fix the help text/description).
2. Thinking models can cause false failures. Grading substring-matches the response, capped at eval_output_tokens. If a thinking model spends its budget in <think> and gets truncated before the answer, a healthy cache scores a fail. 1024 helps but isn't a guarantee — worth a doc note that this budget must exceed the model's thinking allotment.
3. Grader can false-pass. grade_niah_response matches the passkey anywhere in the output, and the prompt itself contains "The secret pass key is {passkey}". A model that echoes prompt text in reasoning but emits garbage as the final answer still passes. Grading the final-answer segment would be more robust.
4. Middle/needle truncation edge case. The safety-net full_tokens = full_tokens[-target_tokens:] in _generate_niah_prompt_text cuts from the front; on a large tokenizer round-trip overshoot with the needle near the 10% depth, the passkey can be clipped out of the prompt entirely → guaranteed (spurious) fail. Rare, but consider trimming suffix filler / preserving the needle span instead.
5. Eval probes count toward perf aggregates. Eval prompts use a 4× output budget and are interleaved into the timed run, so they nudge aggregate output tok/s upward. Since accuracy is tracked separately, consider excluding eval requests from the throughput aggregates.

None of 1–5 block merge — the rebase is the only hard gate. Thanks again!