26dd972adb · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/core/block/cpu_gpu_block_allocator.py

+4 -2

@@ -341,8 +341,10 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):

341 assert device in self._allocators

342 return self._allocators[device].get_prefix_cache_hit_rate()

343

344 - def reset_prefix_cache(self) -> bool:

345 - """Reset prefix cache for all devices."""

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

345 + """Reset prefix cache for specified or all devices."""

346 + if device:

347 + return self._allocators[device].reset_prefix_cache()

348 success = True

349 for allocator in self._allocators.values():

350 success = success and allocator.reset_prefix_cache()

vllm/core/block/interfaces.py

+1 -1

@@ -305,7 +305,7 @@ class DeviceAwareBlockAllocator(ABC):

305 pass

306

307 @abstractmethod

308 - def reset_prefix_cache(self) -> bool:

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

309 """Reset prefix cache."""

310 pass

311

vllm/core/block_manager.py

+2 -2

@@ -456,8 +456,8 @@ class SelfAttnBlockSpaceManager(BlockSpaceManager):

456 def get_prefix_cache_hit_rate(self, device: Device) -> float:

457 return self.block_allocator.get_prefix_cache_hit_rate(device)

458

459 - def reset_prefix_cache(self) -> bool:

460 - return self.block_allocator.reset_prefix_cache()

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

460 + return self.block_allocator.reset_prefix_cache(device)

461

462 def _can_swap(self,

463 seq_group: SequenceGroup,

vllm/core/interfaces.py

+3 -3

@@ -2,7 +2,7 @@

2

3 import enum

4 from abc import ABC, abstractmethod

5 -from typing import List

5 +from typing import List, Optional

6 from typing import Sequence as GenericSequence

7 from typing import Tuple

8

@@ -125,8 +125,8 @@ class BlockSpaceManager(ABC):

125 pass

126

127 @abstractmethod

128 - def reset_prefix_cache(self) -> bool:

129 - """Reset prefix cache for all devices."""

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

129 + """Reset prefix cache for specified or all devices."""

130 pass

131

132 @abstractmethod

vllm/core/placeholder_block_space_manager.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List, Tuple

3 +from typing import List, Optional, Tuple

4

5 from vllm.core.interfaces import AllocStatus, BlockSpaceManager

6 from vllm.sequence import Sequence, SequenceGroup

@@ -92,7 +92,7 @@ class PlaceholderBlockSpaceManager(BlockSpaceManager):

92 def get_prefix_cache_hit_rate(self, device: Device) -> float:

93 return -1

94

95 - def reset_prefix_cache(self) -> bool:

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

96 return True

97

98 def get_num_cached_tokens(self, seq: Sequence) -> int:

vllm/core/scheduler.py

+2 -2

@@ -634,8 +634,8 @@ class Scheduler:

634 def get_prefix_cache_hit_rate(self, device: Device) -> float:

635 return self.block_manager.get_prefix_cache_hit_rate(device)

636

637 - def reset_prefix_cache(self) -> bool:

638 - return self.block_manager.reset_prefix_cache()

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

638 + return self.block_manager.reset_prefix_cache(device)

639

640 def get_num_unfinished_seq_groups(self) -> int:

641 return len(self.waiting) + len(self.running) + len(self.swapped)

vllm/engine/async_llm_engine.py

+4 -3

@@ -35,7 +35,7 @@ from vllm.sampling_params import SamplingParams

35 from vllm.sequence import ExecuteModelRequest

36 from vllm.transformers_utils.tokenizer import AnyTokenizer

37 from vllm.usage.usage_lib import UsageContext

38 -from vllm.utils import deprecate_kwargs, weak_bind

38 +from vllm.utils import Device, deprecate_kwargs, weak_bind

39

40 logger = init_logger(__name__)

41 ENGINE_ITERATION_TIMEOUT_S = envs.VLLM_ENGINE_ITERATION_TIMEOUT_S

@@ -1216,8 +1216,9 @@ class AsyncLLMEngine(EngineClient):

1216 async def stop_profile(self) -> None:

1217 self.engine.stop_profile()

1218

1219 - async def reset_prefix_cache(self) -> None:

1220 - self.engine.reset_prefix_cache()

1219 + async def reset_prefix_cache(self,

1220 + device: Optional[Device] = None) -> None:

1221 + self.engine.reset_prefix_cache(device)

1222

1223 async def sleep(self, level: int = 1) -> None:

1224 self.engine.sleep(level)

vllm/engine/llm_engine.py

+2 -2

@@ -955,12 +955,12 @@ class LLMEngine:

955 """

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

957

958 - def reset_prefix_cache(self) -> bool:

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

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

960

961 success = True

962 for scheduler in self.scheduler:

963 - success = success and scheduler.reset_prefix_cache()

963 + success = success and scheduler.reset_prefix_cache(device)

964 return success

965

966 @staticmethod

vllm/engine/multiprocessing/__init__.py

+4 -3

@@ -13,7 +13,7 @@ from vllm.lora.request import LoRARequest

13 from vllm.outputs import RequestOutput

14 from vllm.prompt_adapter.request import PromptAdapterRequest

15 from vllm.sampling_params import SamplingParams

16 -from vllm.utils import deprecate_kwargs

16 +from vllm.utils import Device, deprecate_kwargs

17

18 VLLM_RPC_SUCCESS_STR = "SUCCESS"

19

@@ -123,8 +123,9 @@ class RPCUProfileRequest(Enum):

123 STOP_PROFILE = 2

124

125

126 -class RPCResetPrefixCacheRequest(Enum):

127 - RESET_PREFIX_CACHE = 1

126 +@dataclass

127 +class RPCResetPrefixCacheRequest:

128 + device: Device

129

130

131 class RPCSleepRequest(Enum):

vllm/engine/multiprocessing/client.py

+4 -3

@@ -47,7 +47,7 @@ from vllm.outputs import PoolingRequestOutput, RequestOutput

47 from vllm.prompt_adapter.request import PromptAdapterRequest

48 from vllm.sampling_params import SamplingParams

49 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs

50 -from vllm.utils import deprecate_kwargs

50 +from vllm.utils import Device, deprecate_kwargs

51

52 logger = init_logger(__name__)

53

@@ -684,11 +684,12 @@ class MQLLMEngineClient(EngineClient):

684 await self._send_one_way_rpc_request(

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

686

687 - async def reset_prefix_cache(self) -> None:

687 + async def reset_prefix_cache(self,

688 + device: Optional[Device] = None) -> None:

689 """Reset the prefix cache"""

690

691 await self._send_one_way_rpc_request(

691 - request=RPCResetPrefixCacheRequest.RESET_PREFIX_CACHE,

692 + request=RPCResetPrefixCacheRequest(device),

693 socket=self.input_socket)

694

695 async def sleep(self, level: int = 1) -> None:

vllm/engine/protocol.py

+3 -2

@@ -18,7 +18,7 @@ from vllm.pooling_params import PoolingParams

18 from vllm.prompt_adapter.request import PromptAdapterRequest

19 from vllm.sampling_params import BeamSearchParams, SamplingParams

20 from vllm.transformers_utils.tokenizer import AnyTokenizer

21 -from vllm.utils import collect_from_async_generator, random_uuid

21 +from vllm.utils import Device, collect_from_async_generator, random_uuid

22

23 logger = init_logger(__name__)

24

@@ -274,7 +274,8 @@ class EngineClient(ABC):

274 ...

275

276 @abstractmethod

277 - async def reset_prefix_cache(self) -> None:

277 + async def reset_prefix_cache(self,

278 + device: Optional[Device] = None) -> None:

279 """Reset the prefix cache"""

280 ...

281

vllm/entrypoints/llm.py

+4 -3

@@ -42,7 +42,8 @@ from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,

42 get_cached_tokenizer)

43 from vllm.transformers_utils.tokenizer_group import TokenizerGroup

44 from vllm.usage.usage_lib import UsageContext

45 -from vllm.utils import Counter, deprecate_args, deprecate_kwargs, is_list_of

45 +from vllm.utils import (Counter, Device, deprecate_args, deprecate_kwargs,

46 + is_list_of)

47

48 logger = init_logger(__name__)

49

@@ -1187,8 +1188,8 @@ class LLM:

1188 def stop_profile(self) -> None:

1189 self.llm_engine.stop_profile()

1190

1190 - def reset_prefix_cache(self) -> bool:

1191 - return self.llm_engine.reset_prefix_cache()

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

1192 + return self.llm_engine.reset_prefix_cache(device)

1193

1194 def sleep(self, level: int = 1):

1195 """

vllm/entrypoints/openai/api_server.py

+7 -3

@@ -85,7 +85,7 @@ from vllm.logger import init_logger

85 from vllm.transformers_utils.config import (

86 maybe_register_config_serialize_by_value)

87 from vllm.usage.usage_lib import UsageContext

88 -from vllm.utils import (FlexibleArgumentParser, get_open_zmq_ipc_path,

88 +from vllm.utils import (Device, FlexibleArgumentParser, get_open_zmq_ipc_path,

89 is_valid_ipv6_address, set_ulimit)

90 from vllm.version import __version__ as VLLM_VERSION

91

@@ -677,8 +677,12 @@ if envs.VLLM_SERVER_DEV_MODE:

677 Reset the prefix cache. Note that we currently do not check if the

678 prefix cache is successfully reset in the API server.

679 """

680 - logger.info("Resetting prefix cache...")

681 - await engine_client(raw_request).reset_prefix_cache()

680 + device = None

681 + device_str = raw_request.query_params.get("device")

682 + if device_str is not None:

683 + device = Device[device_str.upper()]

684 + logger.info("Resetting prefix cache with specific %s...", str(device))

685 + await engine_client(raw_request).reset_prefix_cache(device)

686 return Response(status_code=200)

687

688 @router.post("/sleep")

vllm/v1/engine/async_llm.py

+5 -2

@@ -24,7 +24,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams

24 from vllm.transformers_utils.tokenizer import AnyTokenizer

25 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs

26 from vllm.usage.usage_lib import UsageContext

27 -from vllm.utils import cdiv, kill_process_tree

27 +from vllm.utils import Device, cdiv, kill_process_tree

28 from vllm.v1.engine.core_client import EngineCoreClient

29 from vllm.v1.engine.output_processor import OutputProcessor

30 from vllm.v1.engine.parallel_sampling import ParentRequest

@@ -398,7 +398,10 @@ class AsyncLLM(EngineClient):

398 async def stop_profile(self) -> None:

399 await self.engine_core.profile_async(False)

400

401 - async def reset_prefix_cache(self) -> None:

401 + async def reset_prefix_cache(self,

402 + device: Optional[Device] = None) -> None:

403 + if device == Device.CPU:

404 + raise ValueError("Not supported on CPU.")

405 await self.engine_core.reset_prefix_cache_async()

406

407 async def sleep(self, level: int = 1) -> None:

vllm/v1/engine/llm_engine.py

+2 -1

@@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams

20 from vllm.transformers_utils.tokenizer_group import (

21 BaseTokenizerGroup, init_tokenizer_from_configs)

22 from vllm.usage.usage_lib import UsageContext

23 +from vllm.utils import Device

24 from vllm.v1.engine.core_client import EngineCoreClient

25 from vllm.v1.engine.output_processor import OutputProcessor

26 from vllm.v1.engine.parallel_sampling import ParentRequest

@@ -226,7 +227,7 @@ class LLMEngine:

227 def stop_profile(self):

228 self.engine_core.profile(False)

229

229 - def reset_prefix_cache(self):

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

231 self.engine_core.reset_prefix_cache()

232

233 def sleep(self, level: int = 1):

[FEAT]Support reset prefix cache by specified device (#15003)