@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6cac54f4d1

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 15, 07:15 UTC

[v1] Re-init input batch for multiple kv cache groups (#18654)

Signed-off-by: Chen Zhang <zhangch99@outlook.com>

Chen Zhang committed Jun 4, 2025 at 05:41 UTC 6cac54f4d1673991a415b9897d610c132104155b

6 files changed +61 -46

tests/v1/worker/test_gpu_input_batch.py

+3 -26

@@ -10,8 +10,6 @@ import torch

11 from vllm.sampling_params import SamplingParams

12 from vllm.utils import is_pin_memory_available, make_tensor_with_pad

13 -from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

14 - KVCacheGroupSpec, KVCacheTensor)

13 from vllm.v1.sample.metadata import SamplingMetadata

14 from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable

15 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

@@ -25,27 +23,6 @@ CUDA_DEVICES = [

23 MAX_NUM_PROMPT_TOKENS = 64

28 -def get_kv_cache_config() -> KVCacheConfig:

29 - return KVCacheConfig(

30 - num_blocks=10,

31 - tensors={

32 - "layer.0": KVCacheTensor(size=1024),

33 - },

34 - kv_cache_groups=[

35 - KVCacheGroupSpec(

36 - layer_names=["layer.0"],

37 - kv_cache_spec=FullAttentionSpec(

38 - block_size=1,

39 - num_kv_heads=1,

40 - head_size=16,

41 - dtype=torch.float16,

42 - use_mla=False,

43 - ),

44 - ),

45 - ],

46 - )

47 -

48 -

26 def _compare_objs(obj1, obj2):

27 attrs = inspect.getmembers(obj1, lambda a: not (inspect.isroutine(a)))

28 attr_names = set([

@@ -252,7 +229,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

229 device=torch.device(device),

230 pin_memory=is_pin_memory_available(),

231 vocab_size=1024,

255 - block_size=1,

232 + block_sizes=[1],

233 )

234 reqs: list[CachedRequestState] = []

235 req_id_reqs = {}

@@ -342,7 +319,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,

319 device=torch.device(device),

320 pin_memory=is_pin_memory_available(),

321 vocab_size=1024,

345 - block_size=1,

322 + block_sizes=[1],

323 )

324 ref_input_batch: InputBatch = InputBatch(

325 max_num_reqs=batch_size,

@@ -351,7 +328,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,

328 device=torch.device(device),

329 pin_memory=is_pin_memory_available(),

330 vocab_size=1024,

354 - block_size=1,

331 + block_sizes=[1],

332 )

333

334 reqs: list[CachedRequestState] = []

tests/v1/worker/test_gpu_model_runner.py

+3 -1

@@ -54,7 +54,9 @@ def initialize_kv_cache(runner: GPUModelRunner):

54 device=runner.device,

55 pin_memory=runner.pin_memory,

56 vocab_size=runner.model_config.get_vocab_size(),

57 - block_size=kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size,

57 + block_sizes=[

58 + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size

59 + ],

60 )

61 runner.initialize_attn_backend(kv_cache_config)

vllm/v1/worker/block_table.py

+2 -1

@@ -105,10 +105,11 @@ class MultiGroupBlockTable:

105

106 def __init__(self, max_num_reqs: int, max_model_len: int,

107 max_num_batched_tokens: int, pin_memory: bool,

108 - device: torch.device, block_size: int) -> None:

108 + device: torch.device, block_sizes: list[int]) -> None:

109 self.block_tables = [

110 BlockTable(max_num_reqs, cdiv(max_model_len, block_size),

111 max_num_batched_tokens, pin_memory, device)

112 + for block_size in block_sizes

113 ]

114

115 def append_row(self, block_ids: list[list[int]], row_idx: int) -> None:

vllm/v1/worker/gpu_input_batch.py

+9 -9

@@ -56,14 +56,14 @@ class CachedRequestState:

56 class InputBatch:

58 def __init__(

59 - self,

60 - max_num_reqs: int,

61 - max_model_len: int,

62 - max_num_batched_tokens: int,

63 - device: torch.device,

64 - pin_memory: bool,

65 - vocab_size: int,

66 - block_size: int,

59 + self,

60 + max_num_reqs: int,

61 + max_model_len: int,

62 + max_num_batched_tokens: int,

63 + device: torch.device,

64 + pin_memory: bool,

65 + vocab_size: int,

66 + block_sizes: list[int], # The block_size of each kv cache group

67 ):

68 self.max_num_reqs = max_num_reqs

69 self.max_model_len = max_model_len

@@ -105,7 +105,7 @@ class InputBatch:

105 max_num_batched_tokens=max_num_batched_tokens,

106 pin_memory=pin_memory,

107 device=device,

108 - block_size=block_size,

108 + block_sizes=block_sizes,

109 )

110

111 # Sampling-related.

vllm/v1/worker/gpu_model_runner.py

+40 -6

@@ -143,7 +143,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

143 self.attn_metadata_builders: list[AttentionMetadataBuilder] = []

144 self.attn_backends: list[type[AttentionBackend]] = []

145 # self.kv_cache_config: KVCacheConfig

146 - # self.input_batch: InputBatch # Persistent batch.

146

147 # req_id -> (input_id -> encoder_output)

148 self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}

@@ -173,6 +172,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

172 # Request states.

173 self.requests: dict[str, CachedRequestState] = {}

174

175 + # Input Batch

176 + # NOTE(Chen): Ideally, we should initialize the input batch inside

177 + # `initialize_kv_cache` based on the kv cache config. However, as in

178 + # https://github.com/vllm-project/vllm/pull/18298, due to some unknown

179 + # reasons, we have to initialize the input batch before `load_model`,

180 + # quantization + weight offloading will fail otherwise. As a temporary

181 + # solution, we initialize the input batch here, and re-initialize it

182 + # in `initialize_kv_cache` if the block_sizes here is different from

183 + # the block_sizes in the kv cache config.

184 self.input_batch = InputBatch(

185 max_num_reqs=self.max_num_reqs,

186 max_model_len=self.max_model_len,

@@ -180,7 +188,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

188 device=self.device,

189 pin_memory=self.pin_memory,

190 vocab_size=self.model_config.get_vocab_size(),

183 - block_size=self.cache_config.block_size,

191 + block_sizes=[self.cache_config.block_size],

192 )

193

194 self.use_cuda_graph = (self.vllm_config.compilation_config.level

@@ -2040,6 +2048,35 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2048 self.attn_backends.append(attn_backend_i)

2049 self.attn_metadata_builders.append(attn_metadata_builder_i)

2050

2051 + def may_reinitialize_input_batch(self,

2052 + kv_cache_config: KVCacheConfig) -> None:

2053 + """

2054 + Re-initialize the input batch if the block sizes are different from

2055 + `[self.cache_config.block_size]`. This usually happens when there

2056 + are multiple KV cache groups.

2057 +

2058 + Args:

2059 + kv_cache_config: The KV cache configuration.

2060 + """

2061 + block_sizes = [

2062 + kv_cache_group.kv_cache_spec.block_size

2063 + for kv_cache_group in kv_cache_config.kv_cache_groups

2064 + ]

2065 + if block_sizes != [self.cache_config.block_size]:

2066 + assert self.cache_config.cpu_offload_gb == 0, (

2067 + "Cannot re-initialize the input batch when CPU weight "

2068 + "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501

2069 + "for more details.")

2070 + self.input_batch = InputBatch(

2071 + max_num_reqs=self.max_num_reqs,

2072 + max_model_len=self.max_model_len,

2073 + max_num_batched_tokens=self.max_num_tokens,

2074 + device=self.device,

2075 + pin_memory=self.pin_memory,

2076 + vocab_size=self.model_config.get_vocab_size(),

2077 + block_sizes=block_sizes,

2078 + )

2079 +

2080 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

2081 """

2082 Initialize KV cache based on `kv_cache_config`.

@@ -2047,11 +2084,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2084 kv_cache_config: Configuration for the KV cache, including the KV

2085 cache size of each layer

2086 """

2050 - if len(kv_cache_config.kv_cache_groups) > 1:

2051 - raise NotImplementedError(

2052 - "Hybrid models with more than one KV cache type are not "

2053 - "supported yet.")

2087 self.kv_cache_config = kv_cache_config

2088 + self.may_reinitialize_input_batch(kv_cache_config)

2089 self.initialize_attn_backend(kv_cache_config)

2090

2091 kv_caches: dict[str, torch.Tensor] = {}

vllm/v1/worker/tpu_model_runner.py

+4 -3

@@ -200,7 +200,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):

200 device=self.device,

201 pin_memory=self.pin_memory,

202 vocab_size=self.model_config.get_vocab_size(),

203 - block_size=self.block_size,

203 + block_sizes=[self.block_size],

204 )

205

206 # Cached torch/numpy tensor

@@ -1358,8 +1358,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1358 device=self.device,

1359 pin_memory=self.pin_memory,

1360 vocab_size=self.model_config.get_vocab_size(),

1361 - block_size=kv_cache_config.kv_cache_groups[0].kv_cache_spec.

1362 - block_size,

1361 + block_sizes=[

1362 + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size

1363 + ],

1364 )

1365 # Verify dtype compatibility between block_table_cpu and input_batch

1366 assert self.block_table_cpu.dtype == self.input_batch.block_table[