A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 05:00 UTC
[Bugfix] Align block table for TRTLLM MLA edge-case (#39324)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
May 6, 2026 at 14:17 UTC
38e16678ba7ec8417ce87f7010a03d203ee64b2b
2 files changed
+12
vllm/v1/worker/block_table.py
+7
@@ -257,6 +257,13 @@ class MultiGroupBlockTable:
257
f"must match block_sizes length ({len(block_sizes)})"
258
)
259
260
+ # Align to a multiple of (128 / block_size) as required
261
+ # by some attention backends such as TRTLLM (#39324)
262
+ max_num_blocks = [
263
+ cdiv(n, 128 // bs) * (128 // bs) if bs <= 128 else n
264
+ for n, bs in zip(max_num_blocks, block_sizes)
265
+ ]
266
+
267
self.block_tables = [
268
BlockTable(
269
block_size,
vllm/v1/worker/gpu/block_table.py
+5
@@ -41,6 +41,11 @@ class BlockTables:
41
# As a result, one block on the current rank covers `block_size * cp_size`
42
# tokens in the full, global (unsharded) sequence.
43
max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size)
44
+ # Align to a multiple of (128 / block_size) as required
45
+ # by some attention backends such as TRTLLM (#39324)
46
+ if block_size <= 128:
47
+ alignment = 128 // block_size
48
+ max_num_blocks = cdiv(max_num_blocks, alignment) * alignment
49
block_table = StagedWriteTensor(
50
(self.max_num_reqs, max_num_blocks),
51
dtype=torch.int32,