@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d24cf322e1
Fork 0

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 14, 05:30 UTC

[Hybrid]: Decouple Kernel Block Size from KV Page Size (#24486)

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com> Signed-off-by: Zhiyuan Li <uniartisan2017@gmail.com>

Zhiyuan Li committed Oct 9, 2025 at 14:43 UTC d24cf322e19af19b7a4cbbf108b381b432496f96
18 files changed +572 -54
tests/v1/worker/test_gpu_input_batch.py
+3
@@ -241,6 +241,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
241 pin_memory=is_pin_memory_available(),
242 vocab_size=1024,
243 block_sizes=[1],
244 + kernel_block_sizes=[1],
245 )
246 reqs: list[CachedRequestState] = []
247 req_id_reqs = {}
@@ -335,6 +336,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
336 pin_memory=is_pin_memory_available(),
337 vocab_size=1024,
338 block_sizes=[1],
339 + kernel_block_sizes=[1],
340 )
341 ref_input_batch: InputBatch = InputBatch(
342 max_num_reqs=batch_size,
@@ -344,6 +346,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
346 pin_memory=is_pin_memory_available(),
347 vocab_size=1024,
348 block_sizes=[1],
349 + kernel_block_sizes=[1],
350 )
351
352 reqs: list[CachedRequestState] = []
tests/v1/worker/test_gpu_model_runner.py
+216 -24
@@ -68,6 +68,9 @@ def initialize_kv_cache(runner: GPUModelRunner):
68 pin_memory=runner.pin_memory,
69 vocab_size=runner.model_config.get_vocab_size(),
70 block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
71 + kernel_block_sizes=[
72 + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
73 + ],
74 )
75 runner.initialize_attn_backend(kv_cache_config)
76
@@ -817,42 +820,231 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
820 ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape
821
822 # assert we are using FlashInfer
820 - assert attn_shape[0] == num_blocks
823 + assert attn_shape[0] % num_blocks == 0
824 + block_split_ratio = attn_shape[0] // num_blocks
825 +
826 + # use small blocks for testing to avoid memory issues
827 + test_block_size = min(2, len(blocks0), len(blocks1))
828 +
829 + # use non-overlapping blocks to avoid data contamination
830 + # Split kernel blocks: first half for attention, second half for mamba
831 + mid_point = num_blocks // 2
832 +
833 + # attention uses kernel blocks from first half (mapped to logical blocks)
834 + kv_blocks_for_attention = np.array([0, 1])[:test_block_size]
835 +
836 + # mamba uses kernel blocks from second half
837 + kv_blocks_for_mamba = np.array([mid_point, mid_point + 1])[:test_block_size]
838 +
839 + # create small constant tensors for testing with corrected shapes
840 + # attention: [block_size, ...] starting from dimension 2
841 + attn_constant_shape = attn_shape[2:]
842 + conv_constant_shape = conv_shape[1:]
843 + ssm_constant_shape = ssm_shape[1:]
844
845 attn_blocks_constant = torch.full(
823 - (len(blocks0), *attn_shape[1:]), device=DEVICE, fill_value=3.33
846 + (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33
847 )
848 conv_blocks_constant = torch.full(
826 - (len(blocks1), *conv_shape[1:]), device=DEVICE, fill_value=6.66
849 + (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66
850 )
851 ssm_blocks_constant = torch.full(
829 - (len(blocks1), *ssm_shape[1:]), device=DEVICE, fill_value=9.99
852 + (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99
853 )
854
832 - # fill all attention blocks with constant
855 + # Fill attention blocks with constants using kv block indices
856 + kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio
857 +
858 for layer in [layer_0, layer_1]:
834 - vllm_ctx[layer].kv_cache[0][blocks0, :] = (
835 - attn_blocks_constant.detach().clone()
836 - )
859 + # attention: kv_cache[0][kernel_block_idx, kv_idx, ...]
860 + for i, kernel_block in enumerate(kernel_blocks_for_attention):
861 + vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]
862
838 - # fill all mamba blocks with constant
863 + # fill mamba blocks with constants using kernel block indices
864 for layer in [layer_2, layer_3, layer_4, layer_5]:
840 - vllm_ctx[layer].kv_cache[0][0][blocks1, :] = (
841 - conv_blocks_constant.detach().clone()
842 - )
843 - vllm_ctx[layer].kv_cache[0][1][blocks1, :] = (
844 - ssm_blocks_constant.detach().clone()
845 - )
865 + # mamba: kv_cache[0][component][kernel_block_idx, ...]
866 + for i, kv_block in enumerate(kv_blocks_for_mamba):
867 + vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]
868 + vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]
869
870 # verify attention and mamba contents are correct
871 for layer in [layer_0, layer_1]:
849 - assert torch.equal(
850 - vllm_ctx[layer].kv_cache[0][blocks0, :], attn_blocks_constant
851 - )
872 + for i, kernel_block in enumerate(kernel_blocks_for_attention):
873 + actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]
874 + expected = attn_blocks_constant[i]
875 +
876 + # Check K and V separately
877 + assert torch.equal(actual_kv[0], expected)
878 + assert torch.equal(actual_kv[1], expected)
879 +
880 for layer in [layer_2, layer_3, layer_4, layer_5]:
853 - assert torch.equal(
854 - vllm_ctx[layer].kv_cache[0][0][blocks1, :], conv_blocks_constant
855 - )
856 - assert torch.equal(
857 - vllm_ctx[layer].kv_cache[0][1][blocks1, :], ssm_blocks_constant
858 - )
881 + for i, kv_block in enumerate(kv_blocks_for_mamba):
882 + actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
883 + actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
884 + expected_conv = conv_blocks_constant[i]
885 + expected_ssm = ssm_blocks_constant[i]
886 +
887 + assert torch.equal(actual_conv, expected_conv)
888 + assert torch.equal(actual_ssm, expected_ssm)
889 +
890 + for layer in [layer_2, layer_3, layer_4, layer_5]:
891 + for i, kv_block in enumerate(kv_blocks_for_mamba):
892 + actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
893 + actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
894 + expected_conv = conv_blocks_constant[i]
895 + expected_ssm = ssm_blocks_constant[i]
896 + assert torch.equal(actual_conv, expected_conv)
897 + assert torch.equal(actual_ssm, expected_ssm)
898 +
899 +
900 +def test_hybrid_block_table_initialization():
901 + """Test hybrid block table with different kernel and kvcache_manager block
902 + sizes."""
903 + from vllm.v1.worker.block_table import BlockTable
904 +
905 + # Test configuration: kvcache_manager block size = 32,
906 + # kernel block size = 16
907 + block_size = 32
908 + kernel_block_sizes = [16]
909 + max_num_reqs = 10
910 + max_num_blocks_per_req = 20
911 + max_num_batched_tokens = 512
912 +
913 + block_table = BlockTable(
914 + block_size=block_size,
915 + max_num_reqs=max_num_reqs,
916 + max_num_blocks_per_req=max_num_blocks_per_req,
917 + max_num_batched_tokens=max_num_batched_tokens,
918 + pin_memory=False,
919 + device=torch.device(DEVICE),
920 + kernel_block_size=kernel_block_sizes[0],
921 + )
922 +
923 + # Verify hybrid block configuration
924 + assert block_table.use_hybrid_blocks is True
925 + assert block_table.block_size == kernel_block_sizes[0]
926 + assert block_table.blocks_per_kv_block == (
927 + block_size // kernel_block_sizes[0]
928 + ) # Changed to use first element
929 +
930 + # Test block table conversion logic
931 + # One kvcache_manager block should map to multiple kernel blocks
932 + kvcache_manager_blocks = [0, 1, 2]
933 +
934 + # Verify that kvcache_manager blocks can be converted to kernel blocks
935 + # and that block table operations work correctly.
936 + req_index = 0
937 + block_table.append_row(kvcache_manager_blocks, req_index)
938 + # Get expected kernel blocks from the implementation for verification.
939 + expected_kernel_blocks = block_table._map_to_kernel_blocks(
940 + np.array(kvcache_manager_blocks)
941 + )
942 + # Verify block table state
943 + assert block_table.num_blocks_per_row[req_index] == len(expected_kernel_blocks)
944 + assert np.array_equal(
945 + block_table.block_table.np[req_index, : len(expected_kernel_blocks)],
946 + expected_kernel_blocks,
947 + )
948 +
949 +
950 +def test_input_batch_with_kernel_block_sizes():
951 + """Test InputBatch initialization with kernel_block_sizes parameter."""
952 + max_num_reqs = 10
953 + max_model_len = 512
954 + max_num_batched_tokens = 512
955 + device = torch.device(DEVICE)
956 + pin_memory = False
957 + vocab_size = 50272
958 +
959 + # Test with different kernel block sizes
960 + block_sizes = [32, 64]
961 + kernel_block_sizes = [16, 32]
962 +
963 + input_batch = InputBatch(
964 + max_num_reqs=max_num_reqs,
965 + max_model_len=max_model_len,
966 + max_num_batched_tokens=max_num_batched_tokens,
967 + device=device,
968 + pin_memory=pin_memory,
969 + vocab_size=vocab_size,
970 + block_sizes=block_sizes,
971 + kernel_block_sizes=kernel_block_sizes,
972 + )
973 +
974 + # Verify that block tables were created with kernel block sizes
975 + assert len(input_batch.block_table.block_tables) == len(block_sizes)
976 +
977 + for i, (kv_size, kernel_size) in enumerate(zip(block_sizes, kernel_block_sizes)):
978 + block_table = input_batch.block_table.block_tables[i]
979 + if kv_size != kernel_size:
980 + assert block_table.use_hybrid_blocks is True
981 + assert block_table.block_size == kernel_size
982 + else:
983 + assert block_table.use_hybrid_blocks is False
984 + assert block_table.block_size == kernel_size
985 +
986 +
987 +def test_hybrid_cache_integration(model_runner, dist_init):
988 + """Test hybrid cache architecture integration with GPUModelRunner."""
989 + # Create a new model runner with hybrid cache configuration
990 + vllm_config = get_vllm_config()
991 +
992 + # Configure hybrid cache with different kvcache_manager block size
993 + vllm_config.cache_config.block_size = 32
994 +
995 + model_config = vllm_config.model_config
996 + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
997 + head_size = model_config.get_head_size()
998 + vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
999 + num_heads, head_size, 0.1
1000 + )
1001 +
1002 + runner = GPUModelRunner(vllm_config, DEVICE)
1003 +
1004 + # Initialize KV cache with configuration
1005 + attn_spec = FullAttentionSpec(
1006 + block_size=16, # Use kernel block size directly
1007 + num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config),
1008 + head_size=runner.model_config.get_head_size(),
1009 + dtype=runner.kv_cache_dtype,
1010 + )
1011 + tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
1012 + kv_cache_config = KVCacheConfig(
1013 + num_blocks=NUM_BLOCKS,
1014 + kv_cache_tensors=[
1015 + KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
1016 + ],
1017 + kv_cache_groups=[
1018 + KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
1019 + ],
1020 + )
1021 + runner.kv_cache_config = kv_cache_config
1022 +
1023 + # Initialize input batch with kernel block sizes
1024 + runner.input_batch = InputBatch(
1025 + max_num_reqs=runner.max_num_reqs,
1026 + max_model_len=runner.max_model_len,
1027 + max_num_batched_tokens=runner.max_num_tokens,
1028 + device=runner.device,
1029 + pin_memory=runner.pin_memory,
1030 + vocab_size=runner.model_config.get_vocab_size(),
1031 + block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
1032 + kernel_block_sizes=[16],
1033 + ) # Use kernel block size
1034 +
1035 + runner.initialize_attn_backend(kv_cache_config)
1036 +
1037 + # Verify hybrid block table configuration
1038 + block_table = runner.input_batch.block_table.block_tables[0]
1039 + assert block_table.block_size == (
1040 + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1041 + )
1042 +
1043 + # Test request processing with hybrid blocks
1044 + req_id = "hybrid_req_0"
1045 + scheduler_output = _schedule_new_request(req_id)
1046 +
1047 + # Update states should work with hybrid blocks
1048 + runner._update_states(scheduler_output)
1049 + assert _is_req_scheduled(runner, req_id)
1050 + assert _is_req_state_block_table_match(runner, req_id)
vllm/attention/backends/abstract.py
+17 -1
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 -from typing import Generic, Optional, Protocol, TypeVar
5 +from typing import Generic, Optional, Protocol, TypeVar, Union
6
7 import torch
8
@@ -26,6 +26,13 @@ class AttentionType:
26 """Attention between dec. Q and enc. K/V for encoder-decoder."""
27
28
29 +class MultipleOf:
30 + base: int
31 +
32 + def __init__(self, base: int):
33 + self.base = base
34 +
35 +
36 class AttentionBackend(ABC):
37 """Abstract class for attention backends."""
38
@@ -57,6 +64,10 @@ class AttentionBackend(ABC):
64 def get_metadata_cls() -> type["AttentionMetadata"]:
65 raise NotImplementedError
66
67 + @classmethod
68 + def get_supported_kernel_block_size(cls) -> list[Union[int, MultipleOf]]:
69 + return cls.get_impl_cls().get_supported_kernel_block_size()
70 +
71 @classmethod
72 def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
73 return cls.get_metadata_cls()(*args, **kwargs)
@@ -157,6 +168,11 @@ class AttentionImpl(ABC, Generic[T]):
168 ) -> None:
169 raise NotImplementedError
170
171 + @staticmethod
172 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
173 + # TODO: implement this function for all backends.
174 + return [MultipleOf(1)]
175 +
176 @abstractmethod
177 def forward(
178 self,
vllm/model_executor/models/config.py
+32 -6
@@ -365,6 +365,23 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
365 block_size=model_config.max_model_len,
366 ).page_size_bytes
367
368 + # Model may be marked as is_hybrid
369 + # but mamba is skipped via config,
370 + # return directly
371 + if mamba_page_size == 0:
372 + return
373 +
374 + # Attention backend constraints:
375 + # - FlashAttention (FA) requires block size to be multiple of 16
376 + # - MLA (Multi-head Latent Attention) requires larger alignment:
377 + # * CUTLASS_MLA backend: 128-byte alignment
378 + # * Other MLA backends: 64-byte alignment
379 + if model_config.use_mla:
380 + use_cutlass_mla = envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA"
381 + kernel_block_alignment_size = 128 if use_cutlass_mla else 64
382 + else:
383 + kernel_block_alignment_size = 16
384 +
385 if cache_config.enable_prefix_caching:
386 # With prefix caching, select attention block size to
387 # optimize for mamba kernel performance
@@ -381,19 +398,28 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
398 # TODO(tdoublep): this constraint can be relaxed fairly
399 # easily by changing the way we layout chunks in the
400 # mamba2 kernels.
384 - chunk_size = model_config.get_mamba_chunk_size()
401 +
402 + from math import gcd
403 +
404 + def lcm(a, b):
405 + return a * b // gcd(a, b)
406 +
407 + base_chunk_size = model_config.get_mamba_chunk_size()
408 attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
409 +
410 + chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
411 attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
412 cache_config.mamba_block_size = attn_block_size
413 else:
414 # Without prefix caching, select minimum valid attention block size
415 # to minimize mamba state padding
416
392 - # some attention backends (e.g. FA) only support setting
393 - # block size to multiple of 16, so let's suggest a value
394 - # that would work (note: FA is currently not compatible
395 - # with mamba layers, use FlashInfer instead).
396 - attn_block_size = 16 * cdiv(mamba_page_size, 16 * attn_page_size_1_token)
417 + # Calculate minimum attention block size that satisfies both:
418 + # 1. Backend alignment requirements (kernel_block_alignment_size)
419 + # 2. Mamba page size compatibility (attn_page_size >= mamba_page_size)
420 + attn_block_size = kernel_block_alignment_size * cdiv(
421 + mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token
422 + )
423
424 # override attention block size if either (a) the
425 # user has not set it or (b) the user has set it
vllm/platforms/cuda.py
+19 -7
@@ -118,7 +118,15 @@ class CudaPlatformBase(Platform):
118
119 # TODO(lucas): handle this more gracefully
120 # Note: model_config may be None during testing
121 - if model_config is not None and model_config.use_mla:
121 + # Note: block_size is initialized in
122 + # HybridAttentionMambaModelConfig.verify_and_update_config
123 + # for models with both attention and mamba,
124 + # and doesn't need to be reinitialized here
125 + if (
126 + model_config is not None
127 + and model_config.use_mla
128 + and cache_config.block_size is not None
129 + ):
130 use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
131 # If `VLLM_ATTENTION_BACKEND` is not set and we are using MLA,
132 # then we default to FlashMLA backend for non-blackwell GPUs,
@@ -151,18 +159,22 @@ class CudaPlatformBase(Platform):
159 if (
160 use_flashmla
161 and is_flashmla_dense_supported()[0]
154 - and cache_config.block_size != 64
162 + and cache_config.block_size % 64 != 0
163 ):
164 cache_config.block_size = 64
165 logger.info("Forcing kv cache block size to 64 for FlashMLA backend.")
166
159 - if use_cutlass_mla and cache_config.block_size != 128:
167 + if use_cutlass_mla and cache_config.block_size % 128 != 0:
168 cache_config.block_size = 128
169 logger.info(
170 "Forcing kv cache block size to 128 for CUTLASS_MLA backend."
171 )
172
165 - if use_flashinfer_mla and cache_config.block_size not in [32, 64]:
173 + if (
174 + use_flashinfer_mla
175 + and cache_config.block_size != 32
176 + and cache_config.block_size % 64 != 0
177 + ):
178 cache_config.block_size = 64
179 logger.info(
180 "Forcing kv cache block size to 64 for FlashInferMLA backend."
@@ -269,12 +281,12 @@ class CudaPlatformBase(Platform):
281 use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
282 selected_backend is None
283 and cls.is_device_capability(100)
272 - and block_size == 128
284 + and block_size % 128 == 0
285 )
286 use_flashinfermla = selected_backend == _Backend.FLASHINFER_MLA or (
287 selected_backend is None
288 and cls.is_device_capability(100)
277 - and block_size in [32, 64]
289 + and (block_size == 32 or block_size % 64 == 0)
290 )
291 use_flashmla = selected_backend == _Backend.FLASHMLA or (
292 selected_backend is None and is_flashmla_dense_supported()[0]
@@ -298,7 +310,7 @@ class CudaPlatformBase(Platform):
310 "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
311 )
312 if use_flashmla:
301 - if block_size != 64:
313 + if block_size % 64 != 0:
314 logger.warning(
315 "FlashMLA backend is not supported for block size %d"
316 " (currently only supports block size 64).",
vllm/v1/attention/backends/flash_attn.py
+6 -1
@@ -3,7 +3,7 @@
3 """Attention layer with FlashAttention."""
4
5 from dataclasses import dataclass
6 -from typing import Optional
6 +from typing import Optional, Union
7
8 import numpy as np
9 import torch
@@ -14,6 +14,7 @@ from vllm.attention.backends.abstract import (
14 AttentionImpl,
15 AttentionMetadata,
16 AttentionType,
17 + MultipleOf,
18 is_quantized_kv_cache,
19 )
20 from vllm.attention.layer import Attention
@@ -57,6 +58,10 @@ class FlashAttentionBackend(AttentionBackend):
58 def get_supported_head_sizes(cls) -> list[int]:
59 return [32, 64, 96, 128, 160, 192, 224, 256]
60
61 + @staticmethod
62 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
63 + return [MultipleOf(16)]
64 +
65 @classmethod
66 def validate_head_size(cls, head_size: int) -> None:
67 supported_head_sizes = cls.get_supported_head_sizes()
vllm/v1/attention/backends/flashinfer.py
+8
@@ -23,6 +23,7 @@ from vllm.attention.backends.abstract import (
23 AttentionBackend,
24 AttentionImpl,
25 AttentionType,
26 + MultipleOf,
27 )
28 from vllm.config import CUDAGraphMode, VllmConfig
29 from vllm.logger import init_logger
@@ -165,6 +166,13 @@ class FlashInferBackend(AttentionBackend):
166 # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
167 return [64, 128, 256]
168
169 + @staticmethod
170 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
171 + # Note: Not sure for all platforms,
172 + # but on Blackwell, only support a page size of
173 + # 16, 32, 64
174 + return [16, 32, 64]
175 +
176 @classmethod
177 def validate_head_size(cls, head_size: int) -> None:
178 supported_head_sizes = cls.get_supported_head_sizes()
vllm/v1/attention/backends/mla/cutlass_mla.py
+5
@@ -10,6 +10,7 @@ import vllm._custom_ops as ops
10 from vllm.attention.backends.abstract import (
11 AttentionLayer,
12 AttentionType,
13 + MultipleOf,
14 is_quantized_kv_cache,
15 )
16 from vllm.logger import init_logger
@@ -44,6 +45,10 @@ class CutlassMLABackend(MLACommonBackend):
45 def get_builder_cls() -> type["CutlassMLAMetadataBuilder"]:
46 return CutlassMLAMetadataBuilder
47
48 + @staticmethod
49 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
50 + return [128]
51 +
52
53 class SM100Workspace:
54 def __init__(self, initial_workspace_size):
vllm/v1/attention/backends/mla/flashmla.py
+5 -1
@@ -6,7 +6,7 @@ from typing import ClassVar, Optional, Union
6
7 import torch
8
9 -from vllm.attention.backends.abstract import AttentionLayer, AttentionType
9 +from vllm.attention.backends.abstract import AttentionLayer, AttentionType, MultipleOf
10 from vllm.attention.ops.flashmla import (
11 flash_mla_with_kvcache,
12 get_mla_metadata,
@@ -44,6 +44,10 @@ class FlashMLABackend(MLACommonBackend):
44 def get_impl_cls() -> type["FlashMLAImpl"]:
45 return FlashMLAImpl
46
47 + @staticmethod
48 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
49 + return [64]
50 +
51
52 @dataclass
53 class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
vllm/v1/attention/backends/rocm_aiter_fa.py
+6 -1
@@ -3,7 +3,7 @@
3 """Attention layer with AiterFlashAttention."""
4
5 from dataclasses import dataclass
6 -from typing import Optional
6 +from typing import Optional, Union
7
8 import torch
9
@@ -12,6 +12,7 @@ from vllm.attention.backends.abstract import (
12 AttentionImpl,
13 AttentionMetadata,
14 AttentionType,
15 + MultipleOf,
16 )
17 from vllm.config import VllmConfig
18 from vllm.logger import init_logger
@@ -359,6 +360,10 @@ class AiterFlashAttentionBackend(AttentionBackend):
360 def get_supported_head_sizes(cls) -> list[int]:
361 return [64, 128, 256]
362
363 + @staticmethod
364 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
365 + return [MultipleOf(16)]
366 +
367 @classmethod
368 def validate_head_size(cls, head_size: int) -> None:
369 supported_head_sizes = cls.get_supported_head_sizes()
vllm/v1/attention/backends/tree_attn.py
+6 -1
@@ -4,7 +4,7 @@
4
5 import ast
6 from dataclasses import dataclass
7 -from typing import Optional
7 +from typing import Optional, Union
8
9 import torch
10
@@ -14,6 +14,7 @@ from vllm.attention.backends.abstract import (
14 AttentionImpl,
15 AttentionMetadata,
16 AttentionType,
17 + MultipleOf,
18 )
19 from vllm.attention.ops.triton_unified_attention import unified_attention
20 from vllm.config import VllmConfig
@@ -39,6 +40,10 @@ class TreeAttentionBackend(AttentionBackend):
40 def get_supported_head_sizes(cls) -> list[int]:
41 return [32, 64, 96, 128, 160, 192, 224, 256]
42
43 + @staticmethod
44 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
45 + return [MultipleOf(16)]
46 +
47 @classmethod
48 def validate_head_size(cls, head_size: int) -> None:
49 supported_head_sizes = cls.get_supported_head_sizes()
vllm/v1/attention/backends/triton_attn.py
+6 -1
@@ -3,7 +3,7 @@
3 """High-Performance Triton-only Attention layer."""
4
5 from dataclasses import dataclass
6 -from typing import ClassVar, Optional
6 +from typing import ClassVar, Optional, Union
7
8 import torch
9
@@ -12,6 +12,7 @@ from vllm.attention.backends.abstract import (
12 AttentionImpl,
13 AttentionMetadata,
14 AttentionType,
15 + MultipleOf,
16 )
17 from vllm.attention.ops.triton_reshape_and_cache_flash import (
18 triton_reshape_and_cache_flash,
@@ -157,6 +158,10 @@ class TritonAttentionBackend(AttentionBackend):
158 def get_supported_dtypes(cls) -> list[torch.dtype]:
159 return [torch.float16, torch.bfloat16, torch.float32]
160
161 + @staticmethod
162 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
163 + return [MultipleOf(16)]
164 +
165 @classmethod
166 def validate_head_size(cls, head_size: int) -> None:
167 # Triton Attention supports any head size above 32
vllm/v1/attention/backends/xformers.py
+6 -1
@@ -3,7 +3,7 @@
3 """Attention layer with XFormersAttention."""
4
5 from dataclasses import dataclass
6 -from typing import Optional
6 +from typing import Optional, Union
7
8 import torch
9
@@ -12,6 +12,7 @@ from vllm.attention.backends.abstract import (
12 AttentionImpl,
13 AttentionMetadata,
14 AttentionType,
15 + MultipleOf,
16 )
17 from vllm.attention.ops.triton_unified_attention import unified_attention
18 from vllm.config import VllmConfig
@@ -80,6 +81,10 @@ class XFormersAttentionBackend(AttentionBackend):
81 256,
82 ]
83
84 + @staticmethod
85 + def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
86 + return [MultipleOf(16)]
87 +
88 @classmethod
89 def validate_head_size(cls, head_size: int) -> None:
90 supported_head_sizes = cls.get_supported_head_sizes()
vllm/v1/worker/block_table.py
+85 -4
@@ -22,22 +22,64 @@ class BlockTable:
22 max_num_batched_tokens: int,
23 pin_memory: bool,
24 device: torch.device,
25 + kernel_block_size: int,
26 ):
26 - self.block_size = block_size
27 + """
28 + Args:
29 + block_size: Block size used for KV cache memory allocation
30 + max_num_reqs: Maximum number of concurrent requests supported.
31 + max_num_blocks_per_req: Maximum number of blocks per request.
32 + max_num_batched_tokens: Maximum number of tokens in a batch.
33 + pin_memory: Whether to pin memory for faster GPU transfers.
34 + device: Target device for the block table.
35 + kernel_block_size: The block_size of underlying attention kernel.
36 + Will be the same as `block_size` if `block_size` is supported
37 + by the attention kernel.
38 + """
39 self.max_num_reqs = max_num_reqs
28 - self.max_num_blocks_per_req = max_num_blocks_per_req
40 self.max_num_batched_tokens = max_num_batched_tokens
41 self.pin_memory = pin_memory
42 self.device = device
43
44 + if kernel_block_size == block_size:
45 + # Standard case: allocation and computation use same block size
46 + # No block splitting needed, direct mapping
47 + self.block_size = block_size
48 + self.blocks_per_kv_block = 1
49 + self.use_hybrid_blocks = False
50 + else:
51 + # Hybrid case: allocation block size differs from kernel block size
52 + # Memory blocks are subdivided to match kernel requirements
53 + # Example: 32-token memory blocks with 16-token kernel blocks
54 + # → Each memory block corresponds to 2 kernel blocks
55 + if block_size % kernel_block_size != 0:
56 + raise ValueError(
57 + f"kernel_block_size {kernel_block_size} must divide "
58 + f"kv_manager_block_size size {block_size} evenly"
59 + )
60 +
61 + self.block_size = kernel_block_size
62 + self.blocks_per_kv_block = block_size // kernel_block_size
63 + self.use_hybrid_blocks = True
64 +
65 + self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block
66 +
67 self.block_table = self._make_buffer(
34 - max_num_reqs, max_num_blocks_per_req, dtype=torch.int32
68 + self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32
69 )
70 self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)
71
72 self.slot_mapping = self._make_buffer(
73 self.max_num_batched_tokens, dtype=torch.int64
74 )
75 +
76 + if self.use_hybrid_blocks:
77 + self._kernel_block_arange = np.arange(0, self.blocks_per_kv_block).reshape(
78 + 1, -1
79 + )
80 + else:
81 + self._kernel_block_arange = None
82 +
83 try:
84 self.dcp_world_size = get_dcp_group().world_size
85 self.dcp_rank = get_dcp_group().rank_in_group
@@ -53,6 +95,10 @@ class BlockTable:
95 ) -> None:
96 if not block_ids:
97 return
98 +
99 + if self.use_hybrid_blocks:
100 + block_ids = self._map_to_kernel_blocks(np.array(block_ids))
101 +
102 num_blocks = len(block_ids)
103 start = self.num_blocks_per_row[row_idx]
104 self.num_blocks_per_row[row_idx] += num_blocks
@@ -94,6 +140,7 @@ class BlockTable:
140 req_indices * self.max_num_blocks_per_req
141 + positions // virtual_block_size
142 )
143 +
144 block_numbers = self.block_table.np.ravel()[block_table_indices]
145 # Use virtual_block_size for mask calculation, which marks local
146 # tokens.
@@ -111,6 +158,7 @@ class BlockTable:
158 block_table_indices = (
159 req_indices * self.max_num_blocks_per_req + positions // self.block_size
160 )
161 +
162 block_numbers = self.block_table.np.ravel()[block_table_indices]
163 block_offsets = positions % self.block_size
164 np.add(
@@ -129,6 +177,31 @@ class BlockTable:
177 self.block_table.gpu.fill_(0)
178 self.block_table.cpu.fill_(0)
179
180 + def _map_to_kernel_blocks(self, kv_manager_block_ids: np.ndarray) -> np.ndarray:
181 + """Convert kv_manager_block_id IDs to kernel block IDs.
182 +
183 + Example:
184 + # kv_manager_block_ids: 32 tokens,
185 + # Kernel block size: 16 tokens
186 + # blocks_per_kv_block = 2
187 + >>> kv_manager_block_ids = np.array([0, 1, 2])
188 + >>> Result: [0, 1, 2, 3, 4, 5]
189 +
190 + # Each kv_manager_block_id maps to 2 kernel block id:
191 + # kv_manager_block_id 0 → kernel block id [0, 1]
192 + # kv_manager_block_id 1 → kernel block id [2, 3]
193 + # kv_manager_block_id 2 → kernel block id [4, 5]
194 + """
195 + if not self.use_hybrid_blocks:
196 + return kv_manager_block_ids
197 +
198 + kernel_block_ids = (
199 + kv_manager_block_ids.reshape(-1, 1) * self.blocks_per_kv_block
200 + + self._kernel_block_arange
201 + )
202 +
203 + return kernel_block_ids.reshape(-1)
204 +
205 def get_device_tensor(self, num_reqs: int) -> torch.Tensor:
206 """Returns the device tensor of the block table."""
207 return self.block_table.gpu[:num_reqs]
@@ -160,6 +233,7 @@ class MultiGroupBlockTable:
233 pin_memory: bool,
234 device: torch.device,
235 block_sizes: list[int],
236 + kernel_block_sizes: list[int],
237 num_speculative_tokens: int = 0,
238 ) -> None:
239 # Note(hc): each dcp rank only store
@@ -172,6 +246,12 @@ class MultiGroupBlockTable:
246 # DCP might not be initialized in testing
247 dcp_world_size = 1
248
249 + if len(kernel_block_sizes) != len(block_sizes):
250 + raise ValueError(
251 + f"kernel_block_sizes length ({len(kernel_block_sizes)}) "
252 + f"must match block_sizes length ({len(block_sizes)})"
253 + )
254 +
255 self.block_tables = [
256 BlockTable(
257 block_size,
@@ -183,8 +263,9 @@ class MultiGroupBlockTable:
263 max_num_batched_tokens,
264 pin_memory,
265 device,
266 + kernel_block_size,
267 )
187 - for block_size in block_sizes
268 + for block_size, kernel_block_size in zip(block_sizes, kernel_block_sizes)
269 ]
270
271 def append_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
vllm/v1/worker/gpu_input_batch.py
+2
@@ -78,6 +78,7 @@ class InputBatch:
78 pin_memory: bool,
79 vocab_size: int,
80 block_sizes: list[int], # The block_size of each kv cache group
81 + kernel_block_sizes: list[int],
82 logitsprocs: Optional[LogitsProcessors] = None,
83 is_spec_decode: bool = False,
84 is_pooling_model: bool = False,
@@ -132,6 +133,7 @@ class InputBatch:
133 pin_memory=pin_memory,
134 device=device,
135 block_sizes=block_sizes,
136 + kernel_block_sizes=kernel_block_sizes,
137 num_speculative_tokens=num_speculative_tokens,
138 )
139
vllm/v1/worker/gpu_model_runner.py
+144 -6
@@ -19,7 +19,7 @@ from typing_extensions import TypeAlias
19
20 import vllm.envs as envs
21 from vllm.attention import Attention, AttentionType
22 -from vllm.attention.backends.abstract import AttentionBackend
22 +from vllm.attention.backends.abstract import AttentionBackend, MultipleOf
23 from vllm.attention.layer import MLAAttention
24 from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
25 from vllm.compilation.counter import compilation_counter
@@ -359,6 +359,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
359 pin_memory=self.pin_memory,
360 vocab_size=self.model_config.get_vocab_size(),
361 block_sizes=[self.cache_config.block_size],
362 + kernel_block_sizes=[self.cache_config.block_size],
363 is_spec_decode=bool(self.vllm_config.speculative_config),
364 logitsprocs=build_logitsprocs(
365 self.vllm_config,
@@ -4050,6 +4051,86 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4051 else:
4052 self.reorder_batch_threshold = reorder_batch_threshold_i
4053
4054 + def _find_compatible_block_sizes(
4055 + self,
4056 + kv_manager_block_size: int,
4057 + backend_cls: type[AttentionBackend],
4058 + return_all: bool = False,
4059 + ) -> list[int]:
4060 + """
4061 + Find compatible block sizes for a backend.
4062 +
4063 + Args:
4064 + kv_manager_block_size: Physical block size of KV cache
4065 + backend_cls: Attention backend class
4066 + return_all: Return all compatible sizes if True, max size if False
4067 +
4068 + Returns:
4069 + Compatible block size(s) based on return_all parameter
4070 +
4071 + Raises:
4072 + ValueError: If no compatible block size found
4073 + """
4074 + supported_block_size = backend_cls.get_supported_kernel_block_size()
4075 + compatible_sizes = []
4076 +
4077 + for block_size in supported_block_size:
4078 + if isinstance(block_size, int):
4079 + if kv_manager_block_size % block_size == 0:
4080 + compatible_sizes.append(block_size)
4081 + elif (
4082 + isinstance(block_size, MultipleOf)
4083 + and kv_manager_block_size % block_size.base == 0
4084 + ):
4085 + compatible_sizes.append(kv_manager_block_size)
4086 +
4087 + if not compatible_sizes:
4088 + raise ValueError(f"No compatible block size for {kv_manager_block_size}")
4089 +
4090 + return compatible_sizes if return_all else [max(compatible_sizes)]
4091 +
4092 + def _select_common_block_size(
4093 + self, kv_manager_block_size: int, attn_groups: list[AttentionGroup]
4094 + ) -> int:
4095 + """
4096 + Select common block size for all backends.
4097 +
4098 + Args:
4099 + kv_manager_block_size: Block size of KV cache
4100 + attn_groups: List of attention groups
4101 +
4102 + Returns:
4103 + Block size supported by all backends,
4104 + prioritizing cache_config.block_size
4105 +
4106 + Raises:
4107 + ValueError: If no common block size found
4108 + """
4109 + all_backend_supports = []
4110 +
4111 + for attn_group in attn_groups:
4112 + compatible_sizes = self._find_compatible_block_sizes(
4113 + kv_manager_block_size, attn_group.backend, return_all=True
4114 + )
4115 + supported_sizes = sorted(list(set(compatible_sizes)), reverse=True)
4116 + all_backend_supports.append(set(supported_sizes))
4117 +
4118 + common_supported_sizes = set.intersection(*all_backend_supports)
4119 +
4120 + if not common_supported_sizes:
4121 + error_msg = f"No common block size for {kv_manager_block_size}. "
4122 + for i, attn_group in enumerate(attn_groups):
4123 + supported = all_backend_supports[i]
4124 + error_msg += (
4125 + f"Backend {attn_group.backend} supports: {sorted(supported)}. "
4126 + )
4127 + raise ValueError(error_msg)
4128 +
4129 + if self.cache_config.block_size in common_supported_sizes:
4130 + return self.cache_config.block_size
4131 +
4132 + return max(common_supported_sizes)
4133 +
4134 def may_reinitialize_input_batch(self, kv_cache_config: KVCacheConfig) -> None:
4135 """
4136 Re-initialize the input batch if the block sizes are different from
@@ -4062,8 +4143,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4143 block_sizes = [
4144 kv_cache_group.kv_cache_spec.block_size
4145 for kv_cache_group in kv_cache_config.kv_cache_groups
4146 + if not isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec)
4147 ]
4066 - if block_sizes != [self.cache_config.block_size]:
4148 +
4149 + # Generate kernel_block_sizes that matches each block_size
4150 + kernel_block_sizes = self._prepare_kernel_block_sizes(kv_cache_config)
4151 +
4152 + if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [
4153 + self.cache_config.block_size
4154 + ]:
4155 assert self.cache_config.cpu_offload_gb == 0, (
4156 "Cannot re-initialize the input batch when CPU weight "
4157 "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501
@@ -4077,6 +4165,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4165 pin_memory=self.pin_memory,
4166 vocab_size=self.model_config.get_vocab_size(),
4167 block_sizes=block_sizes,
4168 + kernel_block_sizes=kernel_block_sizes,
4169 is_spec_decode=bool(self.vllm_config.speculative_config),
4170 logitsprocs=self.input_batch.logitsprocs,
4171 is_pooling_model=self.is_pooling_model,
@@ -4128,6 +4217,46 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4217 for attn_groups in self.attn_groups:
4218 yield from attn_groups
4219
4220 + def _prepare_kernel_block_sizes(self, kv_cache_config: KVCacheConfig) -> list[int]:
4221 + """
4222 + Generate kernel_block_sizes that matches each block_size.
4223 +
4224 + For attention backends that support virtual block splitting,
4225 + use the supported block sizes from the backend.
4226 + For other backends (like Mamba), use the same block size (no splitting).
4227 +
4228 + Args:
4229 + kv_cache_config: The KV cache configuration.
4230 +
4231 + Returns:
4232 + list[int]: List of kernel block sizes for each cache group.
4233 + """
4234 + kernel_block_sizes = []
4235 + for kv_cache_group_id, kv_cache_group in enumerate(
4236 + kv_cache_config.kv_cache_groups
4237 + ):
4238 + if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
4239 + continue
4240 + elif isinstance(kv_cache_group.kv_cache_spec, AttentionSpec):
4241 + # This is an attention backend that supports virtual
4242 + # block splitting. Get the supported block sizes from
4243 + # all backends in the group.
4244 + attn_groups = self.attn_groups[kv_cache_group_id]
4245 + kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
4246 + selected_kernel_size = self._select_common_block_size(
4247 + kv_manager_block_size, attn_groups
4248 + )
4249 + kernel_block_sizes.append(selected_kernel_size)
4250 + elif isinstance(kv_cache_group.kv_cache_spec, MambaSpec):
4251 + # This is likely Mamba or other non-attention cache,
4252 + # no splitting.
4253 + kernel_block_sizes.append(kv_cache_group.kv_cache_spec.block_size)
4254 + else:
4255 + raise NotImplementedError(
4256 + f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"
4257 + )
4258 + return kernel_block_sizes
4259 +
4260 def _reshape_kv_cache_tensors(
4261 self,
4262 kv_cache_config: KVCacheConfig,
@@ -4157,16 +4286,24 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4286 num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
4287 if isinstance(kv_cache_spec, AttentionSpec):
4288 has_attn = True
4289 + kv_manager_block_size = kv_cache_spec.block_size
4290 + kernel_size_list = self._find_compatible_block_sizes(
4291 + kv_manager_block_size, attn_backend, return_all=False
4292 + )
4293 + kernel_size = kernel_size_list[0]
4294 + num_blocks_per_kv_block = kv_manager_block_size // kernel_size
4295 + kernel_num_blocks = num_blocks * num_blocks_per_kv_block
4296 +
4297 kv_cache_shape = attn_backend.get_kv_cache_shape(
4161 - num_blocks,
4162 - kv_cache_spec.block_size,
4298 + kernel_num_blocks,
4299 + kernel_size,
4300 kv_cache_spec.num_kv_heads,
4301 kv_cache_spec.head_size,
4302 cache_dtype_str=self.cache_config.cache_dtype,
4303 )
4304 dtype = kv_cache_spec.dtype
4305 try:
4169 - kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
4306 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() # noqa: E501
4307 assert len(kv_cache_stride_order) == len(kv_cache_shape)
4308 except (AttributeError, NotImplementedError):
4309 kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
@@ -4320,10 +4457,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4457 """
4458 kv_cache_config = deepcopy(kv_cache_config)
4459 self.kv_cache_config = kv_cache_config
4323 - self.may_reinitialize_input_batch(kv_cache_config)
4460 self.may_add_encoder_only_layers_to_kv_cache_config()
4461 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
4462 self.initialize_attn_backend(kv_cache_config)
4463 + # Reinitialize need to after initialize_attn_backend
4464 + self.may_reinitialize_input_batch(kv_cache_config)
4465 kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)
4466
4467 if self.speculative_config and self.speculative_config.use_eagle():
vllm/v1/worker/tpu_input_batch.py
+2
@@ -27,6 +27,7 @@ class InputBatch:
27 pin_memory: bool,
28 vocab_size: int,
29 block_sizes: list[int], # The block_size of each kv cache group
30 + kernel_block_sizes: list[int],
31 ):
32 self.max_num_reqs = max_num_reqs
33 self.max_model_len = max_model_len
@@ -68,6 +69,7 @@ class InputBatch:
69 pin_memory=pin_memory,
70 device=device,
71 block_sizes=block_sizes,
72 + kernel_block_sizes=kernel_block_sizes,
73 )
74
75 # Sampling-related.
vllm/v1/worker/tpu_model_runner.py
+4
@@ -259,6 +259,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
259 pin_memory=self.pin_memory,
260 vocab_size=self.model_config.get_vocab_size(),
261 block_sizes=[self.block_size],
262 + kernel_block_sizes=[self.cache_config.block_size],
263 )
264
265 # Cached torch/numpy tensor
@@ -1788,6 +1789,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1789 block_sizes=[
1790 kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1791 ],
1792 + kernel_block_sizes=[
1793 + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1794 + ],
1795 )
1796 # Verify dtype compatibility between block_table_cpu and input_batch
1797 assert (