[V1] Defragmentation support (#1568)

extension PR: HabanaAI/vllm-hpu-extension#275

---------

Signed-off-by: Michal Adamczyk <[email protected]>

Author: madamczyk-intel

requirements/hpu.txt

@@ -7,7 +7,7 @@ ray
 triton==3.1.0
 setuptools>=77.0.3
 setuptools-scm>=8
-vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@1e96318
+vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@cd79204
 
 # Dependencies for HPU vllm docker image
 datasets

vllm/v1/worker/hpu_model_runner.py

@@ -17,11 +17,13 @@
 import torch.distributed
 import vllm_hpu_extension.environment as environment
 from vllm_hpu_extension.bucketing.common import HPUBucketingManager
+from vllm_hpu_extension.defragmentation import OnlineDefragmenter
 from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,
                                          HabanaMemoryProfiler,
                                          HabanaProfilerCounterHelper,
-                                         format_bytes)
-from vllm_hpu_extension.runtime import get_config
+                                         format_bytes, setup_profiler)
+from vllm_hpu_extension.runtime import finalize_config, get_config
+from vllm_hpu_extension.utils import pad_list
 
 from vllm.attention.backends.abstract import AttentionType
 from vllm.attention.layer import Attention
@@ -59,25 +61,6 @@
 _TYPE_CACHE: dict[str, dict[str, Any]] = {}
 
 
-def setup_profiler(warmup, active):
-    schedule = torch.profiler.schedule(wait=0,
-                                       warmup=warmup,
-                                       active=active,
-                                       repeat=1)
-    activities = [
-        torch.profiler.ProfilerActivity.CPU,
-        torch.profiler.ProfilerActivity.HPU
-    ]
-    profiler = torch.profiler.profile(
-        schedule=schedule,
-        activities=activities,
-        on_trace_ready=torch.profiler.tensorboard_trace_handler('.',
-                                                                use_gzip=True),
-        record_shapes=False,
-        with_stack=True)
-    return profiler
-
-
 @dataclass
 class PromptDecodeInfo:
     prompt_req_ids: list[str]
@@ -541,13 +524,6 @@ def round_up(value: int, k: int):
     return (value + k - 1) // k * k
 
 
-def pad_list(input, target_len, val_generator):
-    padding = target_len - len(input)
-    if padding > 0:
-        input.extend(itertools.islice(val_generator, padding))
-    return input
-
-
 class HPUModelRunner:
 
     def __init__(
@@ -558,6 +534,8 @@ def __init__(
     ):
         # TODO: use ModelRunnerBase.__init__(self, vllm_config=vllm_config)
         environment.set_vllm_config(vllm_config)
+        finalize_config()
+
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
         self.cache_config = vllm_config.cache_config
@@ -671,6 +649,8 @@ def __init__(
         self.profiler = HabanaHighLevelProfiler()
         self.profiler_counter_helper = HabanaProfilerCounterHelper()
 
+        self.defragmenter = OnlineDefragmenter()
+
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         """
         Generates the KVCacheSpec by parsing the kv cache format from each
@@ -1075,6 +1055,7 @@ def _extract_prefill_batch_contents(self, num_prefills, num_decodes,
             num_blocks = round_up(context_len + query_len,
                                   self.block_size) // self.block_size
             blocks = block_table_cpu_tensor[batch_idx, :num_blocks].tolist()
+            blocks = [self.defragmenter.resolve(b) for b in blocks]
 
             prompt_tokens = self.input_batch.num_prompt_tokens[batch_idx]
             #TODO: Fix non-prompt case
@@ -1311,6 +1292,8 @@ def _prepare_decode_inputs(self, num_decodes,
                                                   dim=1,
                                                   index=(index //
                                                          self.block_size))
+        block_number.apply_(self.defragmenter.resolve)
+
         block_offsets = padded_index % self.block_size
         slot_mapping = block_number * self.block_size + block_offsets
         # set an out of range value for the padding tokens so that they
@@ -1320,6 +1303,8 @@ def _prepare_decode_inputs(self, num_decodes,
             range(self._PAD_SLOT_ID, self._PAD_SLOT_ID + self.block_size))
         slot_mapping[num_decodes:].apply_(lambda _, ds=dummy_slots: next(ds))
 
+        block_tables_list = self.defragmenter.resolve_all(block_tables_list)
+
         # CONTEXT_LENS [batch_size]
         block_list, block_groups, block_usage = \
             self.get_habana_paged_attn_buffers(
@@ -1598,6 +1583,20 @@ def execute_model(
         # On CPU, sanitize [tokD0, tokD1, tokD2, 0, tokP0, tokP1, tokP2, 0] -> [tokD0, tokD1, tokD2, tokP0, tokP1, tokP2] # noqa
         # Return [tokD0, tokD1, tokD2, tokP0, tokP1, tokP2]
 
+        if self.defragmenter.enabled and self.kv_caches:
+            new = {
+                req.req_id: flatten(req.block_ids)
+                for req in scheduler_output.scheduled_new_reqs if req.block_ids
+            }
+            cached = {
+                req.req_id: flatten(req.new_block_ids)
+                for req in scheduler_output.scheduled_cached_reqs
+                if req.new_block_ids
+            }
+            self.defragmenter.update_state(new | cached,
+                                           scheduler_output.finished_req_ids)
+            self.defragmenter.defragment()
+
         batch_changed = self._update_states(scheduler_output)
         if not scheduler_output.total_num_scheduled_tokens:
             if not has_kv_transfer_group():
@@ -2202,6 +2201,7 @@ def _read_profiling_cfg(self):
 
     @torch.inference_mode()
     def warmup_model(self) -> None:
+        self.defragmenter.initialize(self.kv_caches, self.block_size)
         if not self.enable_bucketing:
             return
         prompt_profile_cfg, decode_profile_cfg = self._read_profiling_cfg()

vllm/v1/worker/hpu_worker.py

@@ -10,7 +10,10 @@
 import torch
 import torch.distributed
 import torch.nn as nn
-from vllm_hpu_extension.profiler import HabanaMemoryProfiler, format_bytes
+from vllm_hpu_extension.debug import init_debug_logger
+from vllm_hpu_extension.profiler import (HabanaMemoryProfiler, format_bytes,
+                                         setup_profiler)
+from vllm_hpu_extension.runtime import get_config
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
@@ -32,6 +35,14 @@
     from vllm.v1.core.scheduler import SchedulerOutput
 
 
+def setup_step_profiler(steps):
+    if steps is None:
+        return None
+    step_start, step_end = steps
+    active = step_end - step_start + 1
+    return setup_profiler(warmup=0, active=active)
+
+
 class HPUWorker:
 
     def __init__(
@@ -76,6 +87,10 @@ def __init__(
         self.gc_track_recompiles = bool(
             "PT_HPU_METRICS_GC_DETAILS" in os.environ
             and bool_helper(os.getenv("PT_HPU_METRICS_GC_DETAILS")))
+        self.step = 0
+        self.profile_steps = get_config().VLLM_PROFILE_STEPS
+        self.step_profiler = setup_step_profiler(self.profile_steps)
+        self.step_debug = init_debug_logger('steps')
 
     def init_profiler(self):
         """Initialize the profiler."""
@@ -254,11 +269,23 @@ def execute_model(
         self,
         scheduler_output: "SchedulerOutput",
     ) -> ModelRunnerOutput:
+        if self.step_debug:
+            self.step_debug(f'step={self.step}')
+        if self.step_profiler and self.step == self.profile_steps[0]:
+            self.step_profiler.start()
         with track_graph_compile('HPUWorker.execute_model') \
             if self.gc_track_recompiles \
             else contextlib.nullcontext():
             output = self.model_runner.execute_model(scheduler_output)
         # TODO(woosuk): Send the output to the engine process.
+        if self.step_profiler:
+            if self.step >= self.profile_steps[0]:
+                self.step_profiler.step()
+            if self.step == self.profile_steps[1]:
+                self.step_profiler.stop()
+                self.step_profiler = None
+                raise RuntimeError('Step profiling finished!')
+        self.step += 1
         return output if self.rank == 0 else None
 
     def profile(self, is_start: bool = True):

vllm/worker/hpu_model_runner.py

@@ -36,7 +36,7 @@
                                          HabanaMemoryProfiler,
                                          HabanaProfilerCounterHelper,
                                          format_bytes)
-from vllm_hpu_extension.runtime import get_config
+from vllm_hpu_extension.runtime import finalize_config, get_config
 
 import vllm.envs as envs
 from vllm.attention import AttentionMetadata, get_attn_backend
@@ -968,7 +968,10 @@ def __init__(
         is_causal: bool = True,
     ):
         ModelRunnerBase.__init__(self, vllm_config=vllm_config)
+
         environment.set_vllm_config(vllm_config)
+        finalize_config()
+
         self.is_driver_worker = is_driver_worker
         self.return_hidden_states = return_hidden_states