feat(specprefill): draft-model-free self-scoring for single-request TTFT

[drumtorben]

Summary

• Adds score_tokens_self() — a SpecPrefill scoring path that works without a draft model, using only the first n_score_layers (default 4) of the target model itself
• Uses an early-exit exception (_ScoringComplete) to stop MLX graph construction after layer N, so Metal only executes ops for those layers (~12.5% of dense prefill cost for a 32-layer model)
• Q vectors from the last n_tail_queries prompt positions act as a proxy for future-generation queries, feeding into the existing _compute_importance → select_chunks → sparse_prefill pipeline unchanged
• Total cost: ~42% of dense prefill (12.5% scoring + 30% sparse prefill at 0.3 keep rate), no second model in RAM required

Motivation

IDE workflows (opencode, Continue, Cursor) suffer a full re-prefill penalty whenever an in-context file changes — the prefix cache ends at [system+tools] and everything after must be re-prefilled densely every turn. Previously, SpecPrefill required a configured draft model (additional 4–8 GB RAM), making it inaccessible to most single-server deployments. This change makes the speedup available to everyone with a single flag.

Architecture safety

No truncation of model.layers is used. make_prompt_cache(model) and _build_layer_to_cache_map() operate on the real model structure throughout, so the following special cases are handled correctly without any extra code:

• Gemma4 (previous_kvs shared-KV cache routing)
• Nemotron-H (block_type sparse attention, compacted layer-to-cache mapping)
• VLM wrappers (VLMModelAdapter.layers property delegation)
• RotatingKVCache (already skipped by _compute_importance)

New API surface

Symbol	Location	Notes
score_tokens_self()	omlx/patches/specprefill.py	Main new function
_ScoringComplete	omlx/patches/specprefill.py	BaseException sentinel
_EarlyExitAttn	omlx/patches/specprefill.py	Installed on layer N to abort graph
_patch_attention_for_capture(..., layers_to_patch=None)	omlx/patches/specprefill.py	Backward-compatible extension
SchedulerConfig.specprefill_self_score_layers	omlx/scheduler.py	Default 4
EngineCore.add_request(..., specprefill_self_score=True)	omlx/engine_core.py	Auto-enables _specprefill_enabled
Request._specprefill_self_score	omlx/request.py	Per-request flag

Test plan

• 9 new unit tests in TestScoreTokensSelf (shape, finiteness, module cleanup on normal + exception path, edge cases for clamped layer/query counts, mx.array input, progress callback, engine propagation)
• All 43 existing test_specprefill.py tests pass
• Full unit test suite: zero new failures (13 pre-existing failures on main unchanged)
• End-to-end: start server, send request with specprefill_self_score=True and prompt >8192 tokens, verify SpecPrefill fires in dashboard and TTFT is reduced

🤖 Generated with Claude Code

[jundot]

Thanks for putting this together. The BaseException early-exit and the clean if/else split from the existing draft path are both nice, easy to read and obviously off-by-default.

Going to close this for now, not because the idea is bad but because I think it needs another pass before it's safe to merge:

1. The flag doesn't reach the HTTP layer. specprefill_self_score is added to EngineCore.add_request() but openai_models.py:252 and the routing block in server.py:2257 are untouched, so no actual /v1/chat/completions caller can turn it on.

2. The "tail prompt Q as a proxy for lookahead Q" substitution is the load-bearing assumption of this PR, and there's no real-model comparison vs the draft-model path. A single run on, say, Llama-3-8B showing selected-index overlap (or that generation quality doesn't regress) would make this much easier to accept.

3. On hybrid-attention models (Gemma-3, gpt-oss-120b style alternating full+sliding layers), the first N attention layers can include RotatingKVCache layers that _compute_importance silently skips. If 2 of 4 layers are sliding the scoring signal is effectively halved, so worth either a smoke test there or an explicit guard.

Happy to re-open and revisit if you want to come back with 1+2 wired up. 3 can be a follow-up. Thanks again for the contribution!