mirror of
https://github.com/TheTom/turboquant_plus.git
synced 2026-10-22 02:53:13 -05:00
c8b5d7ba88
niah_test.py — standalone NIAH benchmark for TurboQuant:
- Spins up llama-server per config, inserts 6-digit "secret code" needles
into filler text, queries via /v1/chat/completions, scores exact match
- Tests all combos of depths x needles x cache types
- Needle positions span 5%-95% of context (not just first half)
- Regex-based scoring (exact 6-digit word boundary match, no false positives)
- Configurable: --depths, --needles, --cache-types, --server-timeout,
--query-timeout, --server-bin, --chars-per-token, --verbose
- Server lifecycle: atexit + signal handlers, stderr captured to temp file
- Deterministic: seed=42, positions logged for reproducibility
- No personal paths in defaults, cross-platform (macOS + Linux)
Hardened per Codex + Roast review:
- Scoring: regex r'\b(\d{6})\b' not substring match (Codex critical)
- Positions: 5%-95% spread not 10%-50% (Roast critical)
- 20 cities (was 10), validates --needles <= city count
- Server timeout 120s (was 60), query timeout 300s (was 120)
- --server-bin flag for ROCm/custom builds
- Server stderr always captured for remote debugging
- Token estimation documented + --chars-per-token override
74 tests, all passing. No /Users/tom/ in any file.
Reference: @jtdavies found TQ 3.5-bit drops to 17-45% of bf16 NIAH
at 128K on MLX. This script lets us measure where we stand on Metal.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: tturney@psyguard.ai