diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index c6f13b8440..03cbd0dc36 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -13,6 +13,7 @@ from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadataBuilder, ) from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.worker.block_table import get_block_table_width def test_indexer_warmup_normalizes_zero_compress_ratios(): @@ -48,11 +49,14 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_ compress_ratio=4, ) vllm_config = create_vllm_config(max_model_len=1024) + max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024) + block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size) builder = DeepseekV32IndexerMetadataBuilder( kv_cache_spec=kv_cache_spec, layer_names=["dummy"], vllm_config=vllm_config, device=device, + block_table_width=block_table_width, ) # Construct a single request where: diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 36ebd1af8f..b90600ce40 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -50,6 +50,7 @@ from vllm.v1.kv_cache_interface import ( KVQuantMode, MLAAttentionSpec, ) +from vllm.v1.worker.block_table import get_block_table_width BACKENDS_TO_TEST = [ AttentionBackendEnum.CUTLASS_MLA, @@ -1495,16 +1496,9 @@ def test_backend_correctness( batch_spec, block_size, device ) - # Pad block table to meet requirement: - # block_num % (128 / block_size) == 0 - required_divisor = int(128 / block_size) current_block_num = common_attn_metadata.block_table_tensor.shape[1] - if current_block_num % required_divisor != 0: - # Pad to next multiple of required_divisor - padded_block_num = ( - (current_block_num + required_divisor - 1) // required_divisor - ) * required_divisor - padding_cols = padded_block_num - current_block_num + padded_block_num = get_block_table_width(current_block_num, block_size) + if padding_cols := padded_block_num - current_block_num: padding = torch.zeros( (common_attn_metadata.block_table_tensor.shape[0], padding_cols), dtype=torch.int32, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index d2d89a74fc..282c022958 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -46,6 +46,11 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.worker.block_table import ( + MultiGroupBlockTable, + SlotMappingMode, + get_block_table_width, +) from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras @@ -1380,6 +1385,29 @@ def test_hybrid_block_table_initialization(): ) +def test_get_block_table_width_aligns_to_128_tokens(): + assert get_block_table_width(1875, 64) == 1876 + + +def test_get_block_table_width_splits_virtual_blocks(): + assert get_block_table_width(235, 256, 64) == 940 + + +def test_mamba_state_table_width_is_not_aligned(): + block_tables = MultiGroupBlockTable( + max_num_reqs=1, + max_num_batched_tokens=1, + pin_memory=False, + device=torch.device("cpu"), + block_sizes=[39664], + kernel_block_sizes=[39664], + max_num_blocks=[1], + slot_mapping_modes=[SlotMappingMode.NONE], + ) + + assert block_tables[0].max_num_blocks_per_req == 1 + + def test_input_batch_with_kernel_block_sizes(): """Test InputBatch initialization with kernel_block_sizes parameter.""" max_num_reqs = 10 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 6f2ddd31f8..97bfd419f8 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -631,6 +631,8 @@ class AttentionMetadataBuilder(ABC, Generic[M]): # Does this backend/builder support updating the block table in existing # metadata supports_update_block_table: bool = False + # Whether the builder constructor requires the block-table width. + requires_block_table_width: ClassVar[bool] = False @abstractmethod def __init__( diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 84692c665f..7051023a7c 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -22,7 +22,6 @@ from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, has_deep_gemm, ) -from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units from vllm.v1.attention.backend import ( AttentionBackend, @@ -37,7 +36,6 @@ from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec -from vllm.v1.worker.cp_utils import get_kv_cache_shard_count logger = init_logger(__name__) @@ -458,6 +456,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # The indexer opts out of the shared reorder-threshold vote (see __init__), # so this is None; its own split uses self.decode_threshold. reorder_batch_threshold: int | None = None + requires_block_table_width = True @classmethod def get_cudagraph_support( @@ -467,7 +466,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) -> AttentionCGSupport: return AttentionCGSupport.UNIFORM_BATCH - def __init__(self, *args, **kwargs): + def __init__(self, *args, block_table_width: int, **kwargs) -> None: super().__init__(*args, **kwargs) scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config @@ -556,15 +555,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): dtype=torch.int32, device=self.device, ) - max_num_blocks_per_req = cdiv( - self.vllm_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_kv_cache_shard_count(), - ) self.expanded_block_table_buffer = torch.zeros( - ( - scheduler_config.max_num_batched_tokens, - max_num_blocks_per_req, - ), + (scheduler_config.max_num_batched_tokens, block_table_width), dtype=torch.int32, device=self.device, ) diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 4d9f256ee0..85228364af 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from enum import Enum import numpy as np @@ -16,6 +17,29 @@ from vllm.v1.utils import CpuGpuBuffer logger = init_logger(__name__) +def get_block_table_width( + max_num_blocks: int, + block_size: int, + kernel_block_size: int | None = None, + *, + token_alignment: int | None = 128, +) -> int: + """Return the width after optional alignment and virtual block splitting.""" + if kernel_block_size is None: + kernel_block_size = block_size + if block_size % kernel_block_size != 0: + raise ValueError( + f"kernel_block_size {kernel_block_size} must divide " + f"block_size {block_size} evenly" + ) + if token_alignment is not None: + if token_alignment <= 0: + raise ValueError("token_alignment must be positive") + block_alignment = token_alignment // math.gcd(token_alignment, block_size) + max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment + return max_num_blocks * block_size // kernel_block_size + + class SlotMappingMode(Enum): TOKEN_TO_KV_SLOT = "token_to_kv_slot" NONE = "none" @@ -277,11 +301,15 @@ 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) + ( + get_block_table_width(n, block_size, token_alignment=None) + if slot_mapping_mode == SlotMappingMode.NONE + else get_block_table_width(n, block_size) + ) + for n, block_size, slot_mapping_mode in zip( + max_num_blocks, block_sizes, slot_mapping_modes + ) ] self.block_tables = [ diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 2183b0df3b..b8e02c289b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -56,6 +56,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput +from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu import pcp_manager as pcp from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput @@ -456,16 +457,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_num_blocks = cdiv( block_table_max_model_len, spec.block_size * self.dcp_size ) - # Align to a multiple of (128 / block_size) as required by some attention - # backends such as TRTLLM (#39324) - if spec.block_size <= 128: - alignment = 128 // spec.block_size - max_num_blocks = cdiv(max_num_blocks, alignment) * alignment # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens if isinstance(spec, MambaSpec): max_num_blocks = ( max_num_blocks if self.cache_config.enable_prefix_caching else 1 ) + spec.num_speculative_blocks + max_num_blocks = get_block_table_width( + max_num_blocks, spec.block_size, token_alignment=None + ) + else: + max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size) max_num_blocks_per_group.append(max_num_blocks) self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 4afb941c12..a8118ef9fa 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -36,6 +36,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.block_table import get_block_table_width logger = init_logger(__name__) @@ -237,12 +238,22 @@ class AttentionGroup: if kernel_block_size is not None else self.kv_cache_spec ) + builder_cls = self.backend.get_builder_cls() + builder_kwargs = {} + if builder_cls.requires_block_table_width: + max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req( + vllm_config, vllm_config.model_config.max_model_len + ) + builder_kwargs["block_table_width"] = get_block_table_width( + max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size + ) self.metadata_builders = [ - self.backend.get_builder_cls()( + builder_cls( kv_cache_spec_builder, self.layer_names, vllm_config, device, + **builder_kwargs, ) for _ in range(num_metadata_builders) ]