[Bugfix] Align block table for TRTLLM MLA edge-case

[benchislett]

Purpose

When running with unusual max-model-len, such as 9416, TRTLLM MLA crashes:

(Worker_TP0 pid=646061) ERROR 04-07 23:47:19 [multiproc_executor.py:932] ValueError: Expected block_num % (128 / block_size) == 0, got block_num=295 and block_size=32`

It seems like a straightforward fix to pad the block table's max_num_blocks by a block or two to avoid this case.

Testing

Ran Kimi K2.5 with the max-model-len and it works.

[gemini-code-assist]

Code Review

This pull request introduces alignment logic for max_num_blocks in vllm/v1/worker/block_table.py and vllm/v1/worker/gpu/block_table.py to meet the requirements of specific attention backends. Feedback was provided regarding the use of magic numbers and code duplication across the two files. Additionally, the reviewer suggested simplifying redundant conditional checks and noted potential issues with the alignment formula if non-power-of-2 block sizes are used.

[mgoin]

Do you think this is something we should add as an attribute to the attention backend? Not sure if this is clearly good to enforce across the board
Cc @LucasWilkinson @MatthewBonanni

[benchislett]

Yeah I'm not sure how to handle this. Seems like a pretty uncommon edge case in the first place so it might not be worth making a big change to support it. But I'm not sure what the end-result impact is to padding the block table by a couple blocks.

[MatthewBonanni]

CUTLASS_MLA has a similar requirement:

vllm/vllm/v1/attention/backends/mla/cutlass_mla.py

Line 203 in 4d042ed

assert block_num % (128 / PAGE_SIZE) == 0

Unless there's significant demand for supporting arbitrary max_num_blocks (when the backend allows) then I'm supportive of always padding for the sake of simplicity

[NickLucche]

@benchislett @LucasWilkinson I guess we're happy with padding the block_table here?

[MatthewBonanni]

LGTM, happy to land once #39324 (comment) is addressed