[Bugfix] Universally align block table width to 128 tokens (#50302)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jul 31, 2026 at 11:52 UTC
a0cd2b69b3dac2b43be02fc16ff940b856d6791b
8 files changed
+90
-30
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+4
@@ -13,6 +13,7 @@ from vllm.v1.attention.backends.mla.indexer import (
13
DeepseekV32IndexerMetadataBuilder,
14
)
15
from vllm.v1.kv_cache_interface import MLAAttentionSpec
16
+from vllm.v1.worker.block_table import get_block_table_width
17
18
19
def test_indexer_warmup_normalizes_zero_compress_ratios():
@@ -48,11 +49,14 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
49
compress_ratio=4,
50
)
51
vllm_config = create_vllm_config(max_model_len=1024)
52
+ max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024)
53
+ block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size)
54
builder = DeepseekV32IndexerMetadataBuilder(
55
kv_cache_spec=kv_cache_spec,
56
layer_names=["dummy"],
57
vllm_config=vllm_config,
58
device=device,
59
+ block_table_width=block_table_width,
60
)
61
62
# Construct a single request where:
tests/v1/attention/test_mla_backends.py
+3
-9
@@ -50,6 +50,7 @@ from vllm.v1.kv_cache_interface import (
50
KVQuantMode,
51
MLAAttentionSpec,
52
)
53
+from vllm.v1.worker.block_table import get_block_table_width
54
55
BACKENDS_TO_TEST = [
56
AttentionBackendEnum.CUTLASS_MLA,
@@ -1495,16 +1496,9 @@ def test_backend_correctness(
1496
batch_spec, block_size, device
1497
)
1498
1498
- # Pad block table to meet requirement:
1499
- # block_num % (128 / block_size) == 0
1500
- required_divisor = int(128 / block_size)
1499
current_block_num = common_attn_metadata.block_table_tensor.shape[1]
1502
- if current_block_num % required_divisor != 0:
1503
- # Pad to next multiple of required_divisor
1504
- padded_block_num = (
1505
- (current_block_num + required_divisor - 1) // required_divisor
1506
- ) * required_divisor
1507
- padding_cols = padded_block_num - current_block_num
1500
+ padded_block_num = get_block_table_width(current_block_num, block_size)
1501
+ if padding_cols := padded_block_num - current_block_num:
1502
padding = torch.zeros(
1503
(common_attn_metadata.block_table_tensor.shape[0], padding_cols),
1504
dtype=torch.int32,
tests/v1/worker/test_gpu_model_runner.py
+28
@@ -46,6 +46,11 @@ from vllm.v1.kv_cache_interface import (
46
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
47
from vllm.v1.sample.metadata import SamplingMetadata
48
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
49
+from vllm.v1.worker.block_table import (
50
+ MultiGroupBlockTable,
51
+ SlotMappingMode,
52
+ get_block_table_width,
53
+)
54
from vllm.v1.worker.gpu.lora_utils import LoraState
55
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
56
from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
@@ -1380,6 +1385,29 @@ def test_hybrid_block_table_initialization():
1385
)
1386
1387
1388
+def test_get_block_table_width_aligns_to_128_tokens():
1389
+ assert get_block_table_width(1875, 64) == 1876
1390
+
1391
+
1392
+def test_get_block_table_width_splits_virtual_blocks():
1393
+ assert get_block_table_width(235, 256, 64) == 940
1394
+
1395
+
1396
+def test_mamba_state_table_width_is_not_aligned():
1397
+ block_tables = MultiGroupBlockTable(
1398
+ max_num_reqs=1,
1399
+ max_num_batched_tokens=1,
1400
+ pin_memory=False,
1401
+ device=torch.device("cpu"),
1402
+ block_sizes=[39664],
1403
+ kernel_block_sizes=[39664],
1404
+ max_num_blocks=[1],
1405
+ slot_mapping_modes=[SlotMappingMode.NONE],
1406
+ )
1407
+
1408
+ assert block_tables[0].max_num_blocks_per_req == 1
1409
+
1410
+
1411
def test_input_batch_with_kernel_block_sizes():
1412
"""Test InputBatch initialization with kernel_block_sizes parameter."""
1413
max_num_reqs = 10
vllm/v1/attention/backend.py
+2
@@ -631,6 +631,8 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
631
# Does this backend/builder support updating the block table in existing
632
# metadata
633
supports_update_block_table: bool = False
634
+ # Whether the builder constructor requires the block-table width.
635
+ requires_block_table_width: ClassVar[bool] = False
636
637
@abstractmethod
638
def __init__(
vllm/v1/attention/backends/mla/indexer.py
+3
-11
@@ -22,7 +22,6 @@ from vllm.utils.deep_gemm import (
22
get_paged_mqa_logits_metadata,
23
has_deep_gemm,
24
)
25
-from vllm.utils.math_utils import cdiv
25
from vllm.utils.platform_utils import num_compute_units
26
from vllm.v1.attention.backend import (
27
AttentionBackend,
@@ -37,7 +36,6 @@ from vllm.v1.attention.backends.utils import (
36
split_decodes_and_prefills,
37
)
38
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
40
-from vllm.v1.worker.cp_utils import get_kv_cache_shard_count
39
40
logger = init_logger(__name__)
41
@@ -458,6 +456,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
456
# The indexer opts out of the shared reorder-threshold vote (see __init__),
457
# so this is None; its own split uses self.decode_threshold.
458
reorder_batch_threshold: int | None = None
459
+ requires_block_table_width = True
460
461
@classmethod
462
def get_cudagraph_support(
@@ -467,7 +466,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
466
) -> AttentionCGSupport:
467
return AttentionCGSupport.UNIFORM_BATCH
468
470
- def __init__(self, *args, **kwargs):
469
+ def __init__(self, *args, block_table_width: int, **kwargs) -> None:
470
super().__init__(*args, **kwargs)
471
scheduler_config = self.vllm_config.scheduler_config
472
parallel_config = self.vllm_config.parallel_config
@@ -556,15 +555,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
555
dtype=torch.int32,
556
device=self.device,
557
)
559
- max_num_blocks_per_req = cdiv(
560
- self.vllm_config.model_config.max_model_len,
561
- self.kv_cache_spec.block_size * get_kv_cache_shard_count(),
562
- )
558
self.expanded_block_table_buffer = torch.zeros(
564
- (
565
- scheduler_config.max_num_batched_tokens,
566
- max_num_blocks_per_req,
567
- ),
559
+ (scheduler_config.max_num_batched_tokens, block_table_width),
560
dtype=torch.int32,
561
device=self.device,
562
)
vllm/v1/worker/block_table.py
+32
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from enum import Enum
6
7
import numpy as np
@@ -16,6 +17,29 @@ from vllm.v1.utils import CpuGpuBuffer
17
logger = init_logger(__name__)
18
19
20
+def get_block_table_width(
21
+ max_num_blocks: int,
22
+ block_size: int,
23
+ kernel_block_size: int | None = None,
24
+ *,
25
+ token_alignment: int | None = 128,
26
+) -> int:
27
+ """Return the width after optional alignment and virtual block splitting."""
28
+ if kernel_block_size is None:
29
+ kernel_block_size = block_size
30
+ if block_size % kernel_block_size != 0:
31
+ raise ValueError(
32
+ f"kernel_block_size {kernel_block_size} must divide "
33
+ f"block_size {block_size} evenly"
34
+ )
35
+ if token_alignment is not None:
36
+ if token_alignment <= 0:
37
+ raise ValueError("token_alignment must be positive")
38
+ block_alignment = token_alignment // math.gcd(token_alignment, block_size)
39
+ max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment
40
+ return max_num_blocks * block_size // kernel_block_size
41
+
42
+
43
class SlotMappingMode(Enum):
44
TOKEN_TO_KV_SLOT = "token_to_kv_slot"
45
NONE = "none"
@@ -277,11 +301,15 @@ class MultiGroupBlockTable:
301
f"must match block_sizes length ({len(block_sizes)})"
302
)
303
280
- # Align to a multiple of (128 / block_size) as required
281
- # by some attention backends such as TRTLLM (#39324)
304
max_num_blocks = [
283
- cdiv(n, 128 // bs) * (128 // bs) if bs <= 128 else n
284
- for n, bs in zip(max_num_blocks, block_sizes)
305
+ (
306
+ get_block_table_width(n, block_size, token_alignment=None)
307
+ if slot_mapping_mode == SlotMappingMode.NONE
308
+ else get_block_table_width(n, block_size)
309
+ )
310
+ for n, block_size, slot_mapping_mode in zip(
311
+ max_num_blocks, block_sizes, slot_mapping_modes
312
+ )
313
]
314
315
self.block_tables = [
vllm/v1/worker/gpu/model_runner.py
+6
-5
@@ -56,6 +56,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
56
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
57
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
58
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
59
+from vllm.v1.worker.block_table import get_block_table_width
60
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
61
from vllm.v1.worker.gpu import pcp_manager as pcp
62
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
@@ -456,16 +457,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
457
max_num_blocks = cdiv(
458
block_table_max_model_len, spec.block_size * self.dcp_size
459
)
459
- # Align to a multiple of (128 / block_size) as required by some attention
460
- # backends such as TRTLLM (#39324)
461
- if spec.block_size <= 128:
462
- alignment = 128 // spec.block_size
463
- max_num_blocks = cdiv(max_num_blocks, alignment) * alignment
460
# For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens
461
if isinstance(spec, MambaSpec):
462
max_num_blocks = (
463
max_num_blocks if self.cache_config.enable_prefix_caching else 1
464
) + spec.num_speculative_blocks
465
+ max_num_blocks = get_block_table_width(
466
+ max_num_blocks, spec.block_size, token_alignment=None
467
+ )
468
+ else:
469
+ max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
470
max_num_blocks_per_group.append(max_num_blocks)
471
472
self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
vllm/v1/worker/utils.py
+12
-1
@@ -36,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
36
MambaSpec,
37
UniformTypeKVCacheSpecs,
38
)
39
+from vllm.v1.worker.block_table import get_block_table_width
40
41
logger = init_logger(__name__)
42
@@ -237,12 +238,22 @@ class AttentionGroup:
238
if kernel_block_size is not None
239
else self.kv_cache_spec
240
)
241
+ builder_cls = self.backend.get_builder_cls()
242
+ builder_kwargs = {}
243
+ if builder_cls.requires_block_table_width:
244
+ max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req(
245
+ vllm_config, vllm_config.model_config.max_model_len
246
+ )
247
+ builder_kwargs["block_table_width"] = get_block_table_width(
248
+ max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size
249
+ )
250
self.metadata_builders = [
241
- self.backend.get_builder_cls()(
251
+ builder_cls(
252
kv_cache_spec_builder,
253
self.layer_names,
254
vllm_config,
255
device,
256
+ **builder_kwargs,
257
)
258
for _ in range(num_metadata_builders)
259
]