[V1][Hybrid] Enable spec decode and optimize block-aligned split in mamba cache align mode

[peakcrosser7]

Purpose

1. Re-enabled spec decoding: Previously disabled in [V1][Hybrid] Mamba Prefix Caching with align mode #30877, speculative decoding is now re-enabled as the related issues have been confirmed as fixed.
2. Optimized block-aligned splitting for resumed requests: Refined the logic to ensure that Mamba states for resumed requests are also cached in a block-aligned fashion, maintaining consistency across prefill phases.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request re-enables speculative decoding for Mamba with align cache mode and refactors the block-aligned splitting logic to better support resumed requests. The changes are logical and well-structured, particularly the introduction of the _mamba_compute_cache_pos helper function. However, I've identified a potential issue in the calculation of the last cacheable position for Eagle mode, which could lead to performance degradation due to cache misses under specific conditions.

[heheda12345]

Is this for correctness or for high cache hit rate?

For cache hit rate of resumed request, I think we only need to ensure the last prefill chunk state was computed before so we have cached it if an external storage exists and don't need to force recompute this position. Try to make the code simple :)

[heheda12345]

add ready for the large unit test

[peakcrosser7]

Is this for correctness or for high cache hit rate?

For cache hit rate of resumed request, I think we only need to ensure the last prefill chunk state was computed before so we have cached it if an external storage exists and don't need to force recompute this position. Try to make the code simple :)

The changes are for both correctness and cache hit rate.
Commit 5781011 is for correctness—ensuring all prefill chunks are block-aligned so the Mamba states are valid.
Commit 1920b70 is the optimization part: this specifically optimizes resumed requests. Previously, prompt and output tokens were treated as a whole for splitting, which often placed the cached state within the output tokens—making it unusable for future request replays. By block-aligned splitting at the original prompt boundary, we cache a state that is much more likely to be hit again. Of course, I didn't consider external storage here.
If you feel this optimization adds too much complexity for such a rare case, I'm fine with removing it to keep the code simple.

[tdoublep]

How will this behave for a normal decode where we have:

num_computed_tokens = request.num_tokens - 1

?

[peakcrosser7]

Oh right, I totally forgot that the normal decode also follows this logic. While the normal decode doesn't get affected by this logic, it still introduces some redundant computations. We could maybe use num_computed_tokens < max(request.num_prompt_tokens, request.num_tokens - 1)? It looks a bit complicated. What do you think? The main point here is just to distinguish the normal decode from the rest, since it doesn't require block-aligned processing.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @peakcrosser7.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[peakcrosser7]

This PR has now been split into #33705 and #33706

[peakcrosser7]

Closed because of #33705 and #33706.