diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index f46e8a8ed63..87a2aac9d4c 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -257,6 +257,13 @@ class MultiGroupBlockTable: f"must match block_sizes length ({len(block_sizes)})" ) + # Align to a multiple of (128 / block_size) as required + # by some attention backends such as TRTLLM (#39324) + max_num_blocks = [ + cdiv(n, 128 // bs) * (128 // bs) if bs <= 128 else n + for n, bs in zip(max_num_blocks, block_sizes) + ] + self.block_tables = [ BlockTable( block_size, diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index d2a2a6aedbf..3061278b019 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -41,6 +41,11 @@ class BlockTables: # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size) + # Align to a multiple of (128 / block_size) as required + # by some attention backends such as TRTLLM (#39324) + if block_size <= 128: + alignment = 128 // block_size + max_num_blocks = cdiv(max_num_blocks, alignment) * alignment block_table = StagedWriteTensor( (self.max_num_reqs, max_num_blocks), dtype=torch.int32,