[CB] Tweaks to update and minor fixes

[remi-or]

Summary

This PR ads minor changes to cache.update, updates the memory handler with all new features and refactors a few parts of the code to make it more readable.

Cache indexing:

• Replace fancy indexing (cache[idx, :, :]) with explicit torch.index_select / index_copy_, which have cleaner behavior under torch.compile and require non-negative indices.
• Switch index storage tensors from int32 to int64 to match index_select/index_copy_ requirements, removing hidden .long() casts in the hot path.
• Introduce sentinel_index and trash_index, which are dedicated positions in the cache padding zone that were already used but now have a name. It also avoids passing negative values (like -1) to indexing functions.

Memory handler (cache.py)

• Collapse the three separate solving methods (compute_num_blocks_and_max_batch_tokens, compute_max_batch_tokens, compute_num_blocks) and verbose compute_memory_footprint into a single polynomial coefficient model. Each term maps to a tensor in _setup_static_tensors, making the memory model auditable and impossible to drift between solvers.
• Account for previously unmodeled tensors: block_table, logprobs output rows, and async double-buffering (when use_async_batching is on).

Benchmark (continuous_batching_overall.py)

• Store results in a timestamped directory instead of a single file, enabling comparison against any previous baseline.

Tests

• Add TestMemoryHandlerPrediction: allocates tensors matching the handler's polynomial model and validates predicted vs actual GPU memory across 5 configurations.
• Fix test_paged_attention: move cache params to ContinuousBatchingConfig, handle list-type eos_token_id, accept attention-impl-dependent output variants.

Performance

Those changes add between 1 and 3% of performance (when not using the block table) depending on the workload. No regressions.

Testing

All tests pass, expect tests/generation/test_continuous_batching.py::ContinuousBatchingWithAcceleratorTest::test_prefix_sharing which is fixed in #45026

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Good job! Nice unbloating!

[ArthurZucker]

very nice!

[remi-or]

thx!

[ArthurZucker]

high school memories