[Performance]: KVCacheManagerV2.get_block_ids_per_seq does per-block Python work on every FlashMLA step #21284

Proposal to improve performance

TrtllmAttentionMetadata.prepare_flash_mla (tensorrt_llm/_torch/attention/backends/trtllm.py) calls kv_cache_manager.get_block_ids_per_seq(self.request_ids) on every forward step. FlashMLA is enabled by default for MLA models with kv_lora_rank + qk_rope_head_dim == 576 (DeepSeek-V3/R1 style) on SM90 with the TRTLLM backend. KVCacheManagerV2 is the default KV cache manager (use_kv_cache_manager_v2=True).

KVCacheManagerV2.get_block_ids_per_seq (tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py) builds its result through get_batch_cache_indices:

1. A Python list comprehension over every block of every request applies the page-index scale.
2. A second comprehension per block applies the layer division and maps BAD_PAGE_INDEX to 0.
3. It builds one torch.tensor per request and pads them with pad_sequence.

So host time per step grows with batch size times blocks per request.

The same file already avoids this in get_batch_cache_indices_flat. It reads each request's page-index buffer with np.frombuffer and applies one vectorized transform. #20848 does the same for the MiniMax-M3 cache manager's own override.

Proposal: build the zero-padded [num_requests, max_blocks] int32 result the same way, with one np.frombuffer copy per row and one vectorized transform. The output (values, shape, dtype) stays the same.

Host-side measurement on CPU, comparing the old and new logic over int32 page buffers shaped like the real ones, with 64-token blocks:

The two versions matched on 300 random cases covering both divisor branches, BAD_PAGE_INDEX entries, and the layer_specs path. I will open a PR with the change and a GPU unit test in tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py. I don't have a GPU, so end-to-end numbers need CI or a maintainer run.

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.