sleep mode level 1

Signed-off-by: Kacper Pietkun <[email protected]>

Author: Kacper-Pietkun

tests/full_tests/sleep_mode.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from vllm import LLM, EngineArgs
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+
+def create_parser():
+    parser = FlexibleArgumentParser()
+    # Add engine args
+    EngineArgs.add_cli_args(parser)
+    parser.set_defaults(model="Qwen/Qwen3-8B", enforce_eager=False)
+    return parser
+
+
+def print_outputs(outputs):
+    print("-" * 50)
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+
+def main(args):
+    """
+    Test script to actually instantiate HPUWorker and test sleep/wakeup functionality.
+    This test creates a real HPUWorker instance and calls the methods.
+    """
+    llm = LLM(**args)
+
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    outputs = llm.generate(prompts)
+    print_outputs(outputs)
+
+    for i in range(3):
+        assert llm.llm_engine.is_sleeping() == False
+        llm.sleep()
+        assert llm.llm_engine.is_sleeping() == True
+        llm.wake_up(["weights"])
+        assert llm.llm_engine.is_sleeping() == True
+        llm.wake_up(["kv_cache"])
+        assert llm.llm_engine.is_sleeping() == False
+        outputs = llm.generate(prompts)
+        print_outputs(outputs)
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    args: dict = vars(parser.parse_args())
+    try:
+        main(args)
+    except Exception:
+        import traceback
+        print("An error occurred during generation:")
+        traceback.print_exc()
+        os._exit(1)

vllm_gaudi/platform.py

@@ -168,6 +168,8 @@ def get_nixl_memory_type(cls) -> str:
             return "VRAM"
         else:
             return "DRAM"
+    def is_sleep_mode_available(cls) -> bool:
+        return True
 
     @classmethod
     def set_torch_compile(cls) -> None:

vllm_gaudi/v1/worker/hpu_worker.py

@@ -10,6 +10,7 @@
 import torch
 import torch.distributed
 import torch.nn as nn
+import habana_frameworks.torch as htorch
 from vllm.tasks import SupportedTask
 from vllm_gaudi.extension.debug import init_debug_logger
 from vllm_gaudi.extension.profiler import (HabanaMemoryProfiler, format_bytes, setup_profiler)
@@ -93,6 +94,10 @@ def __init__(
         self.step_profiler = setup_step_profiler(self.profile_steps)
         self.step_debug = init_debug_logger('steps')
 
+        self.model_sleeping = False
+        self.kv_cache_sleeping = False
+        self.kv_cache_config = None
+
     def init_profiler(self):
         """Initialize the profiler."""
         if envs.VLLM_TORCH_PROFILER_DIR:
@@ -233,6 +238,7 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
         """Allocate GPU KV cache with the specified kv_cache_config."""
 
         with HabanaMemoryProfiler() as m:
+            self.kv_cache_config = kv_cache_config
             self.model_runner.initialize_kv_cache(kv_cache_config)
             torch.hpu.synchronize()
         msg = (f"Usable num_blocks: {kv_cache_config.num_blocks}, "
@@ -316,6 +322,92 @@ def get_kv_connector_handshake_metadata(self) -> dict | None:
         tp_rank = get_tp_group().rank_in_group
         return {tp_rank: metadata}
 
+    def sleep(self, level: int = 1) -> None:
+        """Put the worker into sleep mode to reduce memory usage. Unlike GPU workers that use custom
+        memory allocators, HPU workers use a simpler approach of moving model to CPU and clearing KV cache.
+        Args:
+            level (int): Sleep level (kept for interface compatibility, always performs level 1 operations)
+        """
+
+        assert level == 1, f"Currently, HPU supports only sleep mode level 1 (and not: level {level})"
+        assert not htorch.utils.internal.is_lazy(
+        ) or self.model_config.enforce_eager, "Sleep mode is supported only for torch.compile mode"
+
+        # Handle model - if model was loaded move it to CPU
+        if self.model_sleeping:
+            logger.warning("Model is already in a sleep mode, skipping moving it to CPU")
+        elif not hasattr(self.model_runner, "model") or self.model_runner.model is None:
+            logger.warning("Model was not loaded yet, skipping moving it to CPU")
+        else:
+            with HabanaMemoryProfiler() as m:
+                self.model_runner.model.to("cpu")
+                gc.collect()
+                torch.hpu.synchronize()
+            msg = f"Moving model to CPU for sleep mode took {m.get_summary_string()}"
+            logger.info(msg)
+            self.model_sleeping = True
+
+        # Handle KV cache - discard it
+        if self.kv_cache_sleeping:
+            logger.warning("KV cache has already been discarded by calling sleep method and it has not been reinitialized by calling wake up method yet, skipping discarding it again")
+        elif self.kv_cache_config is None:
+            logger.warning("KV cache has not been initialized yet, skipping discarding it")
+        else:
+            with HabanaMemoryProfiler() as m:
+                self.model_runner.kv_caches = []
+
+                forward_context = self.vllm_config.compilation_config.static_forward_context
+                for layer_name in forward_context:
+                    forward_context[layer_name].kv_cache = None
+
+                gc.collect()
+                torch.hpu.synchronize()
+            msg = f"Discarding KV cache for sleep mode took {m.get_summary_string()}"
+            logger.info(msg)
+            self.kv_cache_sleeping = True
+
+    def wake_up(self, tags: list[str] | None = None) -> None:
+        """Wake up the worker from sleep mode.
+        It can move the model back to HPU and/or reinitialize KV cache.
+        
+        Args:
+            tags: Optional list of tags (kept for interface compatibility)
+        """
+        assert not htorch.utils.internal.is_lazy(
+        ) or self.model_config.enforce_eager, "Sleep mode is supported only for torch.compile mode"
+
+        if tags is None:
+            tags = ["weights", "kv_cache"]
+
+        # Handle model - if model was loaded, move it back to HPU
+        if "weights" in tags:
+            if not self.model_sleeping:
+                logger.warning("Model is not in a sleep mode, skipping moving it to HPU")
+            elif not hasattr(self.model_runner, "model") or self.model_runner.model is None:
+                logger.warning("Model was not loaded yet, skipping moving it to HPU")
+            else:
+                with HabanaMemoryProfiler() as m:
+                    self.model_runner.model.to(self.vllm_config.device_config.device)
+                    gc.collect()
+                    torch.hpu.synchronize()
+                msg = f"Waking up model, moving it back to HPU took {m.get_summary_string()}"
+                logger.info(msg)
+                self.model_sleeping = False
+
+        # Handle KV cache - reinitialize it
+        if "kv_cache" in tags:
+            if not self.kv_cache_sleeping:
+                logger.warning("KV cache is not in a sleep mode, skipping reinitializing it")
+            elif self.kv_cache_config is None:
+                logger.warning("KV cache config is empty, skipping reinitializing KV cache")
+            else:
+                with HabanaMemoryProfiler() as m:
+                    self.model_runner.initialize_kv_cache(self.kv_cache_config)
+                    gc.collect()
+                    torch.hpu.synchronize()
+                msg = f"Waking up KV cache, reinitializing it took {m.get_summary_string()}"
+                logger.info(msg)
+                self.kv_cache_sleeping = False
 
 def init_worker_distributed_environment(
     vllm_config: VllmConfig,
@@ -338,7 +430,6 @@ def init_worker_distributed_environment(
 
 @contextmanager
 def track_graph_compile(name: str):
-    import habana_frameworks.torch as htorch
     from habana_frameworks.torch.hpu.metrics import metric_localcontext
     with metric_localcontext("graph_compilation") as gc:
         yield