61e0a506a3 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/engine/async_llm_engine.py

+3

@@ -1232,6 +1232,9 @@ class AsyncLLMEngine(EngineClient):

1232 async def stop_profile(self) -> None:

1233 self.engine.stop_profile()

1234

1235 + async def reset_mm_cache(self) -> None:

1236 + self.engine.reset_mm_cache()

1237 +

1238 async def reset_prefix_cache(self,

1239 device: Optional[Device] = None) -> None:

1240 self.engine.reset_prefix_cache(device)

vllm/engine/llm_engine.py

+7

@@ -409,6 +409,9 @@ class LLMEngine:

409 # the next step without re-scheduling.

410 self._skip_scheduling_next_step = False

411

412 + # Don't keep the dummy data in memory

413 + self.reset_mm_cache()

414 +

415 def _initialize_kv_caches(self) -> None:

416 """Initialize the KV cache in the worker(s).

417

@@ -913,6 +916,10 @@ class LLMEngine:

916 """

917 return self.scheduler[virtual_engine].has_unfinished_seqs()

918

919 + def reset_mm_cache(self) -> bool:

920 + """Reset the multi-modal cache."""

921 + return self.input_preprocessor.mm_registry.reset_processor_cache()

922 +

923 def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

924 """Reset prefix cache for all devices."""

925

vllm/engine/multiprocessing/__init__.py

+5

@@ -123,6 +123,10 @@ class RPCUProfileRequest(Enum):

123 STOP_PROFILE = 2

124

125

126 +class RPCResetMultiModalCacheRequest(Enum):

127 + RESET = 1

128 +

129 +

130 @dataclass

131 class RPCResetPrefixCacheRequest:

132 device: Device

@@ -164,6 +168,7 @@ class RPCAdapterLoadedResponse:

168

169 RPC_REQUEST_T = Union[RPCProcessRequest, RPCAbortRequest, RPCStartupRequest,

170 RPCUProfileRequest, RPCLoadAdapterRequest,

171 + RPCResetMultiModalCacheRequest,

172 RPCResetPrefixCacheRequest, RPCSleepRequest,

173 RPCWakeUpRequest, RPCIsSleepingRequest]

174

vllm/engine/multiprocessing/client.py

+8

@@ -31,6 +31,7 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

31 RPCIsSleepingResponse,

32 RPCLoadAdapterRequest,

33 RPCProcessRequest,

34 + RPCResetMultiModalCacheRequest,

35 RPCResetPrefixCacheRequest,

36 RPCSleepRequest, RPCStartupRequest,

37 RPCStartupResponse,

@@ -687,6 +688,13 @@ class MQLLMEngineClient(EngineClient):

688 await self._send_one_way_rpc_request(

689 request=RPCUProfileRequest.STOP_PROFILE, socket=self.input_socket)

690

691 + async def reset_mm_cache(self) -> None:

692 + """Reset the multi-modal cache"""

693 +

694 + await self._send_one_way_rpc_request(

695 + request=RPCResetMultiModalCacheRequest.RESET,

696 + socket=self.input_socket)

697 +

698 async def reset_prefix_cache(self,

699 device: Optional[Device] = None) -> None:

700 """Reset the prefix cache"""

vllm/engine/multiprocessing/engine.py

+6

@@ -22,6 +22,7 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

22 RPCIsSleepingResponse,

23 RPCLoadAdapterRequest,

24 RPCProcessRequest,

25 + RPCResetMultiModalCacheRequest,

26 RPCResetPrefixCacheRequest,

27 RPCSleepRequest, RPCStartupRequest,

28 RPCStartupResponse,

@@ -269,6 +270,8 @@ class MQLLMEngine:

270 self.stop_profile()

271 elif isinstance(request, RPCLoadAdapterRequest):

272 self._handle_load_adapter_request(request)

273 + elif isinstance(request, RPCResetMultiModalCacheRequest):

274 + self.reset_mm_cache()

275 elif isinstance(request, RPCResetPrefixCacheRequest):

276 self.reset_prefix_cache()

277 elif isinstance(request, RPCSleepRequest):

@@ -409,6 +412,9 @@ class MQLLMEngine:

412 def stop_profile(self) -> None:

413 self.engine.stop_profile()

414

415 + def reset_mm_cache(self) -> bool:

416 + return self.engine.reset_mm_cache()

417 +

418 def reset_prefix_cache(self) -> bool:

419 return self.engine.reset_prefix_cache()

420

vllm/engine/protocol.py

+5

@@ -278,6 +278,11 @@ class EngineClient(ABC):

278 """Start profiling the engine"""

279 ...

280

281 + @abstractmethod

282 + async def reset_mm_cache(self) -> None:

283 + """Reset the multi-modal cache"""

284 + ...

285 +

286 @abstractmethod

287 async def reset_prefix_cache(self,

288 device: Optional[Device] = None) -> None:

vllm/entrypoints/openai/api_server.py

+4

@@ -150,6 +150,10 @@ async def build_async_engine_client(

150

151 async with build_async_engine_client_from_engine_args(

152 engine_args, args.disable_frontend_multiprocessing) as engine:

153 +

154 + # Don't keep the dummy data in memory

155 + await engine.reset_mm_cache()

156 +

157 yield engine

158

159

vllm/multimodal/processing.py

+5

@@ -1026,6 +1026,11 @@ class ProcessingCache:

1026 def put_item(self, item: ProcessingCacheItem) -> None:

1027 self._cache[item.key] = item.value

1028

1029 + def reset(self) -> bool:

1030 + self._cache.clear()

1031 +

1032 + return True

1033 +

1034

1035 class BaseProcessingInfo:

1036 """Base class to provide the information necessary for data processing."""

vllm/multimodal/registry.py

+10 -4

@@ -88,6 +88,12 @@ class MultiModalRegistry:

88

89 self._processing_cache = ProcessingCache(VLLM_MM_INPUT_CACHE_GIB)

90

91 + def reset_processor_cache(self) -> bool:

92 + """Reset the multi-modal processing cache."""

93 + self._processing_cache.reset()

94 +

95 + return True # Success

96 +

97 @deprecated("Legacy input processor/mapper pipeline has been removed. "

98 "Please update your model runner to use "

99 "`seq_group_metadata.multi_modal_data` directly without "

@@ -106,7 +112,7 @@ class MultiModalRegistry:

112 if not model_config.is_multimodal_model:

113 return {}

114

109 - processor = self.create_processor(model_config, disable_cache=True)

115 + processor = self.create_processor(model_config, disable_cache=False)

116 profiler = MultiModalProfiler(processor)

117

118 seq_len = model_config.max_model_len

@@ -190,7 +196,7 @@ class MultiModalRegistry:

196 if not model_config.is_multimodal_model:

197 return {}

198

193 - processor = self.create_processor(model_config, disable_cache=True)

199 + processor = self.create_processor(model_config, disable_cache=False)

200 profiler = MultiModalProfiler(processor)

201 return profiler.get_mm_limits()

202

@@ -286,7 +292,7 @@ class MultiModalRegistry:

292

293 The model is identified by ``model_config``.

294 """

289 - processor = self.create_processor(model_config, disable_cache=True)

295 + processor = self.create_processor(model_config, disable_cache=False)

296 profiler = MultiModalProfiler(processor)

297 dummy_data = profiler.get_decoder_dummy_data(seq_len, mm_counts)

298

@@ -310,7 +316,7 @@ class MultiModalRegistry:

316

317 The model is identified by ``model_config``.

318 """

313 - processor = self.create_processor(model_config, disable_cache=True)

319 + processor = self.create_processor(model_config, disable_cache=False)

320 profiler = MultiModalProfiler(processor)

321 dummy_data = profiler.get_encoder_dummy_data(seq_len, mm_counts)

322

vllm/v1/engine/async_llm.py

+5

@@ -476,6 +476,11 @@ class AsyncLLM(EngineClient):

476 async def stop_profile(self) -> None:

477 await self.engine_core.profile_async(False)

478

479 + async def reset_mm_cache(self) -> None:

480 + self.processor.mm_registry.reset_processor_cache()

481 + self.processor.mm_input_cache_client.reset()

482 + await self.engine_core.reset_mm_cache_async()

483 +

484 async def reset_prefix_cache(self,

485 device: Optional[Device] = None) -> None:

486 if device == Device.CPU:

vllm/v1/engine/core.py

+9

@@ -286,6 +286,15 @@ class EngineCore:

286 def profile(self, is_start: bool = True):

287 self.model_executor.profile(is_start)

288

289 + def reset_mm_cache(self):

290 + # NOTE: Since this is mainly for debugging, we don't attempt to

291 + # re-sync the internal caches (P0 processor, P0 mirror, P1 mirror)

292 + if self.scheduler.get_num_unfinished_requests():

293 + logger.warning("Resetting the multi-modal cache when requests are "

294 + "in progress may lead to desynced internal caches.")

295 +

296 + self.mm_input_cache_server.reset()

297 +

298 def reset_prefix_cache(self):

299 self.scheduler.reset_prefix_cache()

300

vllm/v1/engine/core_client.py

+15

@@ -88,6 +88,9 @@ class EngineCoreClient(ABC):

88 def profile(self, is_start: bool = True) -> None:

89 raise NotImplementedError

90

91 + def reset_mm_cache(self) -> None:

92 + raise NotImplementedError

93 +

94 def reset_prefix_cache(self) -> None:

95 raise NotImplementedError

96

@@ -143,6 +146,9 @@ class EngineCoreClient(ABC):

146 async def profile_async(self, is_start: bool = True) -> None:

147 raise NotImplementedError

148

149 + async def reset_mm_cache_async(self) -> None:

150 + raise NotImplementedError

151 +

152 async def reset_prefix_cache_async(self) -> None:

153 raise NotImplementedError

154

@@ -214,6 +220,9 @@ class InprocClient(EngineCoreClient):

220 def profile(self, is_start: bool = True) -> None:

221 self.engine_core.profile(is_start)

222

223 + def reset_mm_cache(self) -> None:

224 + self.engine_core.reset_mm_cache()

225 +

226 def reset_prefix_cache(self) -> None:

227 self.engine_core.reset_prefix_cache()

228

@@ -600,6 +609,9 @@ class SyncMPClient(MPClient):

609 def profile(self, is_start: bool = True) -> None:

610 self.call_utility("profile", is_start)

611

612 + def reset_mm_cache(self) -> None:

613 + self.call_utility("reset_mm_cache")

614 +

615 def reset_prefix_cache(self) -> None:

616 self.call_utility("reset_prefix_cache")

617

@@ -787,6 +799,9 @@ class AsyncMPClient(MPClient):

799 async def profile_async(self, is_start: bool = True) -> None:

800 await self.call_utility_async("profile", is_start)

801

802 + async def reset_mm_cache_async(self) -> None:

803 + await self.call_utility_async("reset_mm_cache")

804 +

805 async def reset_prefix_cache_async(self) -> None:

806 await self.call_utility_async("reset_prefix_cache")

807

vllm/v1/engine/llm_engine.py

+8

@@ -101,6 +101,9 @@ class LLMEngine:

101 # for v0 compatibility

102 self.model_executor = self.engine_core.engine_core.model_executor # type: ignore

103

104 + # Don't keep the dummy data in memory

105 + self.reset_mm_cache()

106 +

107 @classmethod

108 def from_vllm_config(

109 cls,

@@ -240,6 +243,11 @@ class LLMEngine:

243 def stop_profile(self):

244 self.engine_core.profile(False)

245

246 + def reset_mm_cache(self):

247 + self.processor.mm_registry.reset_processor_cache()

248 + self.processor.mm_input_cache_client.reset()

249 + self.engine_core.reset_mm_cache()

250 +

251 def reset_prefix_cache(self, device: Optional[Device] = None):

252 self.engine_core.reset_prefix_cache()

253

vllm/v1/engine/mm_input_cache.py

+5

@@ -83,3 +83,8 @@ class MirroredProcessingCache:

83 full_mm_inputs.append(mm_input)

84

85 return full_mm_inputs

86 +

87 + def reset(self) -> bool:

88 + self.mm_cache.clear()

89 +

90 + return True

vllm/v1/engine/processor.py

+4

@@ -54,6 +54,10 @@ class Processor:

54 self.use_hash = self.mm_input_cache_client.use_cache or \

55 self.cache_config.enable_prefix_caching

56

57 + @property

58 + def mm_registry(self):

59 + return self.input_preprocessor.mm_registry

60 +

61 def _validate_logprobs(

62 self,

63 params: SamplingParams,

[Bugfix] Avoid repeatedly creating dummy data during engine startup (#17935)