Automatically optimize transformer loglikelihood scoring

[Qubitium]

Summary

Adds scorer-side prefix KV reuse for repeated one-token multiple-choice loglikelihood workloads (MMLU), while preserving the existing exact A/B/C/D grouping.

The scorer now:

• groups question contexts by their shared fixed prefix;
• prefills each eligible subject prefix once;
• batches divergent question suffixes;
• gathers A/B/C/D probabilities from one final-logit row per question;
• reports warmup misses separately from steady-state cache hits;
• batches distinct prefix prewarms (default cap 32) and releases completed MMLU subject-group entries.

Prefix caching and prewarming are enabled by default for transformer loglikelihood sessions; callers can still opt out with loglikelihood_prefix_cache=False or --no-loglikelihood-prefix-cache-prewarm.

H200 validation

Llama 3.2 1B QVQ checkpoint, MMLU humanities, 32 rows, batch size 128, paged|sdpa:

Revision	Score phase	Total task	Accuracy
Pre-cache Evalution d966bee	71.805 s	109.876 s	0.3438
This PR	12.751 s	50.438 s	0.3438

This is 5.63x faster in the scoring phase with identical accuracy. The current run recorded one compulsory warmup miss followed by 1/1 steady-state scoring hits (100%).

The new H200 128-row lifecycle run (0c74736) records 51.308 s of scoring and 89.916 s total, versus 51.416 s and 91.541 s for the same-head no-prewarm control. It achieved 5/5 steady-state hits, used 2 batched warmup forwards for 5 prefixes, and finished with 0 cache entries after 2 release calls (accuracy 0.2812).

Commit bd8ba27 moves lifecycle release into the generic transformer loglikelihood and loglikelihood_continuous paths, so non-MMLU loglikelihood suites receive the same bounded cache and end-of-call cleanup automatically. MMLU retains subject-contiguous submission to release each subject's entries promptly.

Review follow-ups d13c302 and 3219460 make batched cache-row extraction device-local per layer for sharded/offloaded caches, fall back to independent prefill for unknown cache wrappers, cap prewarm work to LRU capacity, remove dead helper code, and propagate OOM/illegal-access/fatal CUDA errors instead of hiding them behind an uncached retry.

Tests

• 25 focused transformer/MMLU loglikelihood tests passed (73 passed, 3 skipped across both files).
• Version remains 0.0.17, exactly one patch above the latest PyPI release (0.0.16); subsequent commits do not bump it.

[Qubitium]

Reviewed together with QVQ #83. The ownership split is correct: Evalution should own the scorer/cache lifecycle and QVQ should only forward overrides. The one-token scoring math also looks exact for causal models: score each candidate from the final context-token logits and keep the divergent suffix path on a copied prefix cache.

I have two issues I would fix before merging because this is now default-on for transformer loglikelihood:

1. BLOCKER — batched prewarm row selection is not safe for sharded/device-mapped KV caches. _select_loglikelihood_cache_row() creates indices on self.input_device and passes that single tensor to cache.batch_select_indices(indices) (and uses it for legacy index_select). HF Cache.batch_select_indices() delegates the same index tensor to every cache layer. With device_map/model sharding/offload, KV layers can live on cuda:0, cuda:1, CPU, etc., so an index tensor pinned to the input device is not valid for every cache tensor. The first multi-prefix prewarm can therefore raise on a large multi-GPU model and then disable the optimization for the session. The H200 single-GPU validation cannot exercise this. Please make row selection device-agnostic/per-layer (or fall back to unbatched prewarm when the cache is distributed) and add a unit test that exercises a cache whose rows are not all on input_device.

2. BLOCKER — do not treat every RuntimeError as a cache compatibility failure. Both _score_one_token_choice_prefix_cache() and _prewarm_loglikelihood_prefix_cache_for_batch() catch RuntimeError, disable the cache, and fall back. That also catches torch.OutOfMemoryError, CUDA illegal-access/runtime failures, and unrelated model bugs. With a default-on optimization this can hide the real failure and immediately retry a larger uncached/full-context forward after an OOM. Restrict fallback to known cache-API incompatibilities (TypeError/AttributeError/NotImplementedError, or narrowly identified runtime compatibility errors) and re-raise OOM/CUDA/fatal runtime errors.

Non-blocking performance note: if one scoring batch has more eligible prefixes than loglikelihood_prefix_cache_max_entries, _prewarm_loglikelihood_prefix_cache_for_batch() prewarms all chunks before scoring; later chunks can evict earlier warm entries before they are consumed. Consider capping prewarm to retained capacity or interleaving prewarm+score by capacity-sized groups.

Otherwise the generic loglikelihood / loglikelihood_continuous finally cleanup and MMLU subject-contiguous lifecycle are the right direction.

[Qubitium]

Review follow-up complete.

• d13c302 performs per-layer, tensor-device-local KV row selection for sharded/offloaded caches and falls back to unbatched prefill for unknown cache wrappers.
• d13c302 restricts RuntimeError fallback to cache-compatibility messages; OOM, illegal access, device asserts, launch failures, and other fatal CUDA errors propagate.
• Prewarm work is capped to the configured LRU capacity.
• 3219460 removes the obsolete batch-size helper.

Validation: 73 passed, 3 skipped in tests/test_transformer.py tests/test_mmlu.py.