diff --git a/tests/v1/attention/test_indexer_expanded_block_table.py b/tests/v1/attention/test_indexer_expanded_block_table.py new file mode 100644 index 000000000000..b348f6ebe930 --- /dev/null +++ b/tests/v1/attention/test_indexer_expanded_block_table.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the DSA indexer's expanded block table.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import create_vllm_config +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, +) +from vllm.v1.kv_cache_interface import ( + MLAAttentionSpec, + get_block_table_width, +) +from vllm.v1.worker.block_table import MultiGroupBlockTable + + +def _make_builder(block_table_width: int, max_num_batched_tokens: int = 16): + builder = object.__new__(DeepseekV32IndexerMetadataBuilder) + builder.device = torch.device("cpu") + builder.expanded_block_table_buffer = torch.zeros( + (max_num_batched_tokens, block_table_width), dtype=torch.int32 + ) + builder.decode_seq_lens_buffer = torch.zeros( + max_num_batched_tokens, dtype=torch.int32 + ) + builder.arange_buffer = torch.arange(max_num_batched_tokens, dtype=torch.int32) + builder.decode_lens_buffer = torch.zeros(max_num_batched_tokens, dtype=torch.int32) + return builder + + +def test_nonuniform_decode_uses_finalized_block_table_width(): + block_tables = MultiGroupBlockTable( + max_num_reqs=2, + max_num_batched_tokens=8, + pin_memory=False, + device=torch.device("cpu"), + block_sizes=[64], + kernel_block_sizes=[64], + max_num_blocks=[1875], + ) + block_table = block_tables[0].get_device_tensor(2) + assert block_table.shape == (2, 1876) + indexer_width = get_block_table_width(1875, 64, 64) + assert indexer_width == block_table.shape[1] + builder = _make_builder(indexer_width) + block_table.copy_(torch.arange(2 * 1876, dtype=torch.int32).view(2, 1876)) + decode_lens_cpu = torch.tensor([4, 2], dtype=torch.int32) + + _, expanded_block_table, _, _, _ = builder._prepare_decode_tensors( + seq_lens=torch.tensor([100, 100], dtype=torch.int32), + block_table=block_table, + decode_lens=decode_lens_cpu, + decode_lens_cpu=decode_lens_cpu, + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + num_decodes=2, + num_decode_tokens=8, + use_native=False, + next_n=4, + max_decode_len=4, + ) + + expected = torch.repeat_interleave(block_table, decode_lens_cpu, dim=0) + torch.testing.assert_close(expanded_block_table[:6], expected) + assert expanded_block_table.shape == (8, 1876) + + +def test_block_table_width_aligns_before_kernel_block_splitting(): + block_tables = MultiGroupBlockTable( + max_num_reqs=1, + max_num_batched_tokens=1, + pin_memory=False, + device=torch.device("cpu"), + block_sizes=[256], + kernel_block_sizes=[64], + max_num_blocks=[235], + ) + + expected_width = get_block_table_width(235, 256, 64) + assert expected_width == 940 + assert block_tables[0].get_device_tensor(1).shape[1] == expected_width + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_indexer_buffer_accounts_for_dcp_and_kernel_block_splitting(monkeypatch): + kv_cache_block_size = 256 + kernel_block_size = 64 + vllm_config = create_vllm_config(max_model_len=1200, block_size=kv_cache_block_size) + vllm_config.parallel_config.decode_context_parallel_size = 2 + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.indexer.get_dcp_group", + lambda: SimpleNamespace(rank_in_group=0), + ) + kv_cache_spec = MLAAttentionSpec( + block_size=kv_cache_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + ).copy_with_new_block_size(kernel_block_size) + + builder = DeepseekV32IndexerMetadataBuilder( + kv_cache_spec=kv_cache_spec, + layer_names=["dummy"], + vllm_config=vllm_config, + device=torch.device("cuda"), + ) + + max_num_kv_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1200) + expected_width = get_block_table_width( + max_num_kv_blocks, kv_cache_block_size, kernel_block_size + ) + assert expected_width == 12 + assert builder.expanded_block_table_buffer.shape[1] == expected_width diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 31acd475adec..2896fb6468bf 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -19,7 +19,7 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch): block_sizes=[16, 32, 8], max_num_reqs=4, max_num_batched_tokens=64, - max_num_blocks_per_group=[8, 8, 8], + block_table_widths=[8, 16, 8], device=device, kernel_block_sizes=[16, 16, 8], ) @@ -113,7 +113,7 @@ def test_block_tables_apply_staged_writes_single_group(): block_sizes=[16], max_num_reqs=2, max_num_batched_tokens=16, - max_num_blocks_per_group=[4], + block_table_widths=[4], device=device, kernel_block_sizes=[16], ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 79e3a60e9810..be8caacb33e1 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1329,6 +1329,7 @@ def test_hybrid_attention_mamba_tensor_shapes(): def test_hybrid_block_table_initialization(): """Test hybrid block table with different kernel and kvcache_manager block sizes.""" + from vllm.v1.kv_cache_interface import get_block_table_width from vllm.v1.worker.block_table import BlockTable # Test configuration: kvcache_manager block size = 32, @@ -1343,7 +1344,9 @@ def test_hybrid_block_table_initialization(): block_table = BlockTable( block_size=block_size, max_num_reqs=max_num_reqs, - max_num_blocks_per_req=max_num_blocks_per_req, + block_table_width=get_block_table_width( + max_num_blocks_per_req, block_size, kernel_block_sizes[0] + ), max_num_batched_tokens=max_num_batched_tokens, pin_memory=False, device=torch.device(DEVICE_TYPE), diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 8cbfa539b7ed..c20201b7b472 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -226,7 +226,7 @@ def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: block_table = BlockTable( block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, max_num_reqs=1, - max_num_blocks_per_req=block_table_stride, + block_table_width=block_table_stride, max_num_batched_tokens=num_tokens, pin_memory=False, device=device, diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py index 8d2328432ebc..5cec50a39c4e 100644 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -29,7 +29,7 @@ def warm_v1_block_table_kernels( block_table = BlockTable( block_size=block_size, max_num_reqs=1, - max_num_blocks_per_req=max_num_blocks_per_req, + block_table_width=max_num_blocks_per_req, max_num_batched_tokens=max(num_tokens, max_tokens), pin_memory=False, device=device, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index fe8bfef53088..76eae365cdb5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -22,7 +22,6 @@ 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, @@ -36,8 +35,11 @@ get_dcp_local_seq_lens, split_decodes_and_prefills, ) -from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec -from vllm.v1.worker.cp_utils import get_kv_cache_shard_count +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + MLAAttentionSpec, + get_block_table_width, +) logger = init_logger(__name__) @@ -556,15 +558,16 @@ def __init__(self, *args, **kwargs): 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(), + max_num_blocks_per_req = self.kv_cache_spec.max_num_blocks_per_req( + self.vllm_config, self.vllm_config.model_config.max_model_len + ) + max_num_blocks_per_req = get_block_table_width( + max_num_blocks_per_req, + self.kv_cache_spec.block_table_block_size, + self.kv_cache_spec.block_size, ) self.expanded_block_table_buffer = torch.zeros( - ( - scheduler_config.max_num_batched_tokens, - max_num_blocks_per_req, - ), + (scheduler_config.max_num_batched_tokens, max_num_blocks_per_req), dtype=torch.int32, device=self.device, ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 164350c7fb17..22b87ddc6710 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -5,7 +5,7 @@ import copy from collections import Counter -from dataclasses import dataclass, fields, replace +from dataclasses import dataclass, field, fields, replace from enum import Enum, IntEnum from math import prod from typing import TYPE_CHECKING @@ -83,6 +83,27 @@ def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: return get_kv_quant_mode(kv_cache_dtype).is_per_token_head +def align_block_table_width(max_num_blocks: int, block_size: int) -> int: + """Align a block-table row to the width required by attention backends.""" + if block_size > 128: + return max_num_blocks + alignment = 128 // block_size + return cdiv(max_num_blocks, alignment) * alignment + + +def get_block_table_width( + max_num_blocks: int, block_size: int, kernel_block_size: int +) -> int: + """Return the block-table width after alignment and block splitting.""" + if block_size % kernel_block_size != 0: + raise ValueError( + f"kernel_block_size {kernel_block_size} must divide " + f"block_size {block_size} evenly" + ) + max_num_blocks = align_block_table_width(max_num_blocks, block_size) + return max_num_blocks * block_size // kernel_block_size + + class KVCacheSpecKind(str, Enum): FULL_ATTENTION = "full_attention" MLA_ATTENTION = "mla_attention" @@ -177,6 +198,7 @@ class AttentionSpec(KVCacheSpec): num_kv_heads: int head_size: int dtype: torch.dtype + kv_cache_block_size: int | None = field(default=None, compare=False, repr=False) kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None indexes_kv_by_block_stride: bool = False @@ -220,7 +242,18 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config kv_shard_count = parallel_config.decode_context_parallel_size - return cdiv(max_len, self.block_size * kv_shard_count) + return cdiv(max_len, self.block_table_block_size * kv_shard_count) + + @property + def block_table_block_size(self) -> int: + return self.kv_cache_block_size or self.block_size + + def copy_with_new_block_size(self, block_size: int) -> Self: + return replace( + self, + block_size=block_size, + kv_cache_block_size=self.block_table_block_size, + ) @dataclass(frozen=True, kw_only=True) diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 332eda4cbdfe..1b3554aa9839 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -9,8 +9,8 @@ from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import get_block_table_width from vllm.v1.utils import CpuGpuBuffer logger = init_logger(__name__) @@ -26,7 +26,7 @@ def __init__( self, block_size: int, max_num_reqs: int, - max_num_blocks_per_req: int, + block_table_width: int, max_num_batched_tokens: int, pin_memory: bool, device: torch.device, @@ -38,7 +38,7 @@ def __init__( Args: block_size: Block size used for KV cache memory allocation max_num_reqs: Maximum number of concurrent requests supported. - max_num_blocks_per_req: Maximum number of blocks per request. + block_table_width: Number of entries in each block-table row. max_num_batched_tokens: Maximum number of tokens in a batch. pin_memory: Whether to pin memory for faster GPU transfers. device: Target device for the block table. @@ -76,7 +76,7 @@ def __init__( self.blocks_per_kv_block = block_size // kernel_block_size self.use_hybrid_blocks = True - self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block + self.max_num_blocks_per_req = block_table_width self.block_table = self._make_buffer( self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32 @@ -272,18 +272,16 @@ def __init__( 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) + block_table_widths = [ + get_block_table_width(n, bs, kbs) + for n, bs, kbs in zip(max_num_blocks, block_sizes, kernel_block_sizes) ] self.block_tables = [ BlockTable( block_size, max_num_reqs, - max_num_blocks_per_req, + block_table_width, max_num_batched_tokens, pin_memory, device, @@ -294,10 +292,13 @@ def __init__( for ( block_size, kernel_block_size, - max_num_blocks_per_req, + block_table_width, slot_mapping_mode, ) in zip( - block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes + block_sizes, + kernel_block_sizes, + block_table_widths, + slot_mapping_modes, ) ] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 22c4afc11bc1..d60e03b3dcf9 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -20,7 +20,7 @@ def __init__( block_sizes: list[int], max_num_reqs: int, max_num_batched_tokens: int, - max_num_blocks_per_group: list[int], + block_table_widths: list[int], device: torch.device, kernel_block_sizes: list[int], cp_size: int = 1, @@ -38,7 +38,7 @@ def __init__( self.cp_interleave = cp_interleave self.num_kv_cache_groups = len(self.block_sizes) - assert len(max_num_blocks_per_group) == self.num_kv_cache_groups + assert len(block_table_widths) == self.num_kv_cache_groups self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -47,9 +47,10 @@ def __init__( # num_kv_cache_groups x [max_num_reqs, max_num_blocks] self.block_tables: list[StagedWriteTensor] = [] for i in range(self.num_kv_cache_groups): - max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i] block_table = StagedWriteTensor( - (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device + (self.max_num_reqs, block_table_widths[i]), + dtype=torch.int32, + device=device, ) self.block_tables.append(block_table) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fdedbfb86d03..e8b9e130a13e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -54,7 +54,12 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib 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.kv_cache_interface import ( + KVCacheConfig, + MambaSpec, + align_block_table_width, + get_block_table_width, +) from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu import pcp_manager as pcp @@ -444,13 +449,11 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: 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 = align_block_table_width( + max_num_blocks, spec.block_size + ) max_num_blocks = ( max_num_blocks if self.cache_config.enable_prefix_caching else 1 ) + spec.num_speculative_blocks @@ -459,11 +462,24 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) + block_table_widths = [] + for group_id, (kv_cache_group, max_num_blocks) in enumerate( + zip(kv_cache_config.kv_cache_groups, max_num_blocks_per_group) + ): + spec = kv_cache_group.kv_cache_spec + if not isinstance(spec, MambaSpec): + max_num_blocks = get_block_table_width( + max_num_blocks, + spec.block_size, + self.kernel_block_sizes[group_id], + ) + block_table_widths.append(max_num_blocks) + self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, max_num_batched_tokens=self.max_num_tokens, - max_num_blocks_per_group=max_num_blocks_per_group, + block_table_widths=block_table_widths, device=self.device, kernel_block_sizes=self.kernel_block_sizes, cp_size=self.dcp_size,