[Feature]: Support kv cache's enable_block_reuse feature in AutoDeploy's integration of flashinfer

[suyoggupta]

🚀 The feature, motivation and pitch

Currently, setting enable_block_reuse to true in kv_cache_config will trigger an IMA when flashinfer is used as the attention backend in AutoDeploy

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[lucaslie]

This is partially supported via #8108. Specifically, #8108 adds support for enable_block_reuse=True, enable_partial_reuse=True, copy_on_partial_reuse=False.

Follow-up items:

• copy_on_partial_reuse=True --> AD only uses the "page assignment" functionality of the KVCacheManager while maintaining its own cache pool. copy_on_partial_reuse enables copying partial non-leaf blocks into a new block to be used by a new request. We currently do not support the functionality of doing partial copying to a new block id (see 3.ii. below)
• reliably disable block reuse for SSMs --> since SSM state is per sequence we cannot reuse blocks and must disable it by default. Why? --> The state always reflects the whole history of the sequence. You can never "roll it back" or truncate it on a per-token basis like the kv-cache. You also cannot reuse blocks in hybrid models since the SSM layers need the entire hidden state with all context tokens to recompute their state. Moreover, there is no way in TRT-LLM at the moment to represent state across multiple requests to indicate a continuation request. All requests are independent and chat continuation for regular attention-only models is simply represented via block reuse, see slack discussion https://nvidia.slack.com/archives/C059LSY62BT/p1759352324194129

Some notes on the KVCacheManager

1. blocks are represented in a radix search tree
2. full blocks are reused by simply pointing to the block for a new request. Since full blocks are read-only they can safely be reused multiple times and by multiple requests simultaneously
3. partial blocks can be reused under the following circumstances:
   1. if they are a leaf block (i.e. correspond to the terminal block of a sequence) and are not in use anymore, they can be safely reclaimed by a new sequence and simply "rolled" back to the part of the block that is a match with the new sequence and then used from that position onwards
   2. if they are a non-leaf block (i.e. a full block that is part of a longer sequence with subsequent blocks) they cannot be safely re-used or reclaimed. However, the parts of the block that are matching can be copied over the a new block and used by the new sequence

Currently, AD supports (2.) and (3.i.) but not (3.ii.)

[suyoggupta]

done in #8108