[Bugfix] Size the sparse-indexer expanded block table from the runner's block-table width

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

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],
     )

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),

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,

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,

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,
         )

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)

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,
             )
         ]
 

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)