fix(scheduler): clear SpecPrefill state after failures and cache rejects - #4035

Problem

If a SpecPrefill request fails after its sparse prefill, stale SpecPrefill state can stop the engine from admitting requests, while stale sparse indices can make fallback paths prefill the wrong tokens.

• A decode error that reaches fail_all_requests (from the engine loop's exception handler) leaves _specprefill_active_request_id and the offset RoPE patch in place with no request behind them. _schedule_waiting defers everything from then on.
• A cache-corruption error goes through _recover_from_cache_error and puts the request back in waiting, but the id stays set, so the retried request is deferred by its own leftover id and never finishes.
• _reset_request_for_reprefill re-prefills from the full prompt with cached_tokens = 0, but scoring doesn't run again. After a prefix-cache hit the request keeps indices that point into the old suffix, so the retry prefills the wrong tokens at the wrong positions without any error.
• _validate_cache has the same problem. Scoring runs during the prefix-cache lookup, before validation, so when validation rejects the restored cache the fallback switches to the full prompt but keeps the suffix-relative indices. The regression test covers the recurrent-slot-count rejection path using a hybrid model and a persistent cache block whose recurrent state doesn't match the current model.

Related: #2336. Its Bug B is the _recover_from_cache_error case above, reported there as a hang with TurboQuant KV and SpecPrefill on a real model. Bug A (the sparse-prefilled cache not being TurboQuant-converted) is not touched here.

Fix

• fail_all_requests and _recover_from_cache_error call _cleanup_specprefill, as abort, normal completion and generation-overflow recovery already do.
• _reset_request_for_reprefill drops the indices, so the retry prefills densely, like the existing fallback when sparse prefill fails. This also applies to the generation-overflow retry, which uses the same reset.
• The _validate_cache rejection fallback drops the indices scored against the rejected cache (second commit).

Regression coverage

Tests reach each state through EngineCore.generate with tiny random-init models and the repo's mock tokenizer. The only faults are a model forward that raises once during decode and, for the _validate_cache case, a stored block rewritten on the SSD cache between two engine sessions.

On main the three test cases in test_specprefill_lifecycle.py time out with the engine stuck, and in the _validate_cache test the model is given a sparse prefill of suffix tokens where the full prompt should start. Removing either index reset on the branch makes its test fail: the corruption retry, or the _validate_cache fallback, feeds the model suffix tokens instead of the full prompt.

Validation

• tests/test_specprefill_lifecycle.py and tests/test_specprefill_cache_validation.py: 4 passed on the branch, 4 failed on main
• tests/test_specprefill*.py: 111 passed
• tests/test_scheduler.py tests/test_engine_core.py tests/test_prefill_oom_graceful.py tests/test_exceptions.py: 623 passed
• python -m pytest: 14984 passed, 381 skipped, 86 deselected, 0 failed
• ruff check and black --check on the two new test files. omlx/scheduler.py already doesn't pass Black on main; the lines added here do.

Not tested

I haven't run a representative real-model hardware check for this change. The tests above use tiny random-init models, including the hybrid cache-validation case, so they exercise the scheduler and persistent-cache paths but not a production-size model or the TurboQuant workload from #2336.