Port high-level profiler to V1 engine (#1501)

Author: jkaniecki

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@d928f25
+vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@cbe274b
 
 # Dependencies for HPU vllm docker image
 datasets

vllm/v1/worker/hpu_model_runner.py

@@ -16,7 +16,10 @@
 import torch.distributed
 import vllm_hpu_extension.environment as environment
 from vllm_hpu_extension.bucketing.common import HPUBucketingManager
-from vllm_hpu_extension.profiler import HabanaMemoryProfiler, format_bytes
+from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,
+                                         HabanaMemoryProfiler,
+                                         HabanaProfilerCounterHelper,
+                                         format_bytes)
 from vllm_hpu_extension.runtime import get_config
 
 from vllm.attention.backends.abstract import AttentionType
@@ -525,6 +528,7 @@ def __init__(
         self,
         vllm_config: VllmConfig,
         device: torch.device = 'hpu',
+        is_driver_worker: bool = False,
     ):
         # TODO: use ModelRunnerBase.__init__(self, vllm_config=vllm_config)
         environment.set_vllm_config(vllm_config)
@@ -538,6 +542,7 @@ def __init__(
         self.speculative_config = vllm_config.speculative_config
         self.prompt_adapter_config = vllm_config.prompt_adapter_config
         self.observability_config = vllm_config.observability_config
+        self.is_driver_worker = is_driver_worker
 
         self.sampler = get_sampler()
 
@@ -636,6 +641,9 @@ def __init__(
         # TODO(madamczyk-intel): add a knob for that
         # TODO(madamczyk-intel): debug why increasing it lowers acc
         self.logits_rounding = 1
+        # High-level profiler
+        self.profiler = HabanaHighLevelProfiler()
+        self.profiler_counter_helper = HabanaProfilerCounterHelper()
 
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         """
@@ -841,6 +849,7 @@ def _get_prompts_and_decodes(
 
         # Traverse decodes first
         decode_req_ids = []
+        num_computed_tokens_decode = []
         for i in range(num_reqs):
             req_id = self.input_batch.req_ids[i]
             assert req_id is not None
@@ -857,6 +866,11 @@ def _get_prompts_and_decodes(
             # This is decode
             assert num_scheduled_tokens == 1
             decode_req_ids.append(req_id)
+            num_computed_tokens_decode.append(int(num_computed_tokens + 1))
+
+        if self.profiler.enabled:
+            self.profiler_counter_helper.capture_decode_seq_stats(
+                num_computed_tokens_decode)
 
         # Traverse prompts
         prompt_req_ids = []
@@ -1071,6 +1085,8 @@ def _form_prefill_batch(self, contents):
         token_ids = contents.token_ids
         req_ids = contents.req_ids
         query_lens = [len(tids) for tids in contents.token_ids]
+        if self.profiler.enabled:
+            self.profiler_counter_helper.capture_prompt_seq_stats(query_lens)
         context_lens = contents.context_lens
 
         token_positions = [
@@ -1375,17 +1391,31 @@ def _execute_model_generic(self,
             # no hpu graphs for t.compile?
             use_graphs = False
         trimmed_attn_metadata = trim_attn_metadata(attn_metadata)
-        hidden_states = self.model.forward(input_ids=token_ids,
-                                           positions=position_ids,
-                                           attn_metadata=trimmed_attn_metadata,
-                                           kv_caches=kv_caches)
+        if self.is_driver_worker:
+            model_event_name = ("model_forward_"
+                                f"bs{batch_size}_"
+                                f"seq{seq_len}_"
+                                f"ctx{num_blocks}_"
+                                f"graphs{'T' if use_graphs else 'F'}")
+        else:
+            model_event_name = 'model_executable'
+        with self.profiler.record_event('internal', model_event_name):
+            hidden_states = self.model.forward(
+                input_ids=token_ids,
+                positions=position_ids,
+                attn_metadata=trimmed_attn_metadata,
+                kv_caches=kv_caches)
         # NOTE(kzawora): returning hidden_states is required in prompt logprobs
         # scenarios, as they will do logit processing on their own
         non_flattened_hidden_states = hidden_states
 
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
         hidden_states = hidden_states[logits_indices]
-        logits = self.model.compute_logits(hidden_states, None)
+        with self.profiler.record_event('internal', ('compute_logits'
+                                                     f'{batch_size}_'
+                                                     f'seq{seq_len}_ctx'
+                                                     f'{num_blocks}')):
+            logits = self.model.compute_logits(hidden_states, None)
         return non_flattened_hidden_states, logits
 
     def _get_prompt_logprobs_dict(
@@ -1532,8 +1562,9 @@ def execute_model(
         num_decodes = len(pd_info.decode_req_ids)
         num_prefills = len(pd_info.prompt_req_ids)
         num_reqs = num_decodes + num_prefills
-        prefill_data, decode_data = self._prepare_inputs(
-            scheduler_output, num_prefills, num_decodes)
+        with self.profiler.record_event('internal', 'prepare_input_tensors'):
+            prefill_data, decode_data = self._prepare_inputs(
+                scheduler_output, num_prefills, num_decodes)
 
         #FIXME(kzawora): Currently there's no handling of logprobs. Fix that
         # later.
@@ -1548,64 +1579,102 @@ def execute_model(
                       attn_metadata, logits_indices,
                       logits_requests) in enumerate(
                           zip(*shallow_tuple(prefill_data))):
+                self.event_start = self.profiler.get_timestamp_us()
+                self.profiler.start("internal", "prefill")
                 htorch.core.mark_step()
                 prefill_hidden_states_ts, logits_device = \
                     self._execute_model_generic(
                         token_ids, position_ids, attn_metadata, logits_indices,
                         self.kv_caches)
                 htorch.core.mark_step()
-                sampling_metadata = self._prepare_sampling(
-                    batch_changed, req_id, pad_to=logits_device.shape[0])
-                sampler_output = self.sampler(
-                    logits=logits_device, sampling_metadata=sampling_metadata)
-                prefill_sampled_token_ids.append(
-                    sampler_output.sampled_token_ids.flatten())
-                prefill_sampled_requests.extend(logits_requests)
+                with self.profiler.record_event('internal', "sampler"):
+                    sampling_metadata = self._prepare_sampling(
+                        batch_changed, req_id, pad_to=logits_device.shape[0])
+                    sampler_output = self.sampler(
+                        logits=logits_device,
+                        sampling_metadata=sampling_metadata)
+                    prefill_sampled_token_ids.append(
+                        sampler_output.sampled_token_ids.flatten())
+                    prefill_sampled_requests.extend(logits_requests)
                 htorch.core.mark_step()
+                if self.is_driver_worker and self.profiler.enabled:
+                    # Stop recording 'execute_model_generic' event
+                    self.profiler.end()
+                    event_end = self.profiler.get_timestamp_us()
+                    counters = self.profiler_counter_helper.get_counter_dict(
+                        cache_config=self.cache_config,
+                        duration=event_end - self.event_start,
+                        seq_len=self._seq_len(attn_metadata),
+                        batch_size_padded=token_ids.size(0),
+                        real_batch_size=len(req_id),
+                        prompt_batch_idx=idx,
+                        is_prompt=True)
+                    self.profiler.record_counter(self.event_start, counters)
+            if self.is_driver_worker and self.profiler.enabled:
+                self.profiler_counter_helper.reset_prompt_seq_stats()
 
         ######################### DECODES #########################
         # Decodes run as one single batch with [padded_decode_bs, 1]
         if num_decodes > 0:
+            self.event_start = self.profiler.get_timestamp_us()
+            self.profiler.start("internal", "decode")
             assert decode_data is not None
             htorch.core.mark_step()
             _, logits_device = self._execute_model_generic(
                 decode_data.token_ids, decode_data.position_ids,
                 decode_data.attn_metadata, decode_data.logits_indices,
                 self.kv_caches)
             htorch.core.mark_step()
-            sampling_metadata = self._prepare_sampling(
-                batch_changed,
-                pd_info.decode_req_ids,
-                pad_to=logits_device.shape[0])
-            sampler_output = self.sampler(logits=logits_device,
-                                          sampling_metadata=sampling_metadata)
-            decode_sampled_token_ids.append(
-                sampler_output.sampled_token_ids.flatten())
-            decode_sampled_requests.extend(
-                self.input_batch.req_ids[:num_decodes])
+            with self.profiler.record_event('internal', "sampler"):
+                sampling_metadata = self._prepare_sampling(
+                    batch_changed,
+                    pd_info.decode_req_ids,
+                    pad_to=logits_device.shape[0])
+                sampler_output = self.sampler(
+                    logits=logits_device, sampling_metadata=sampling_metadata)
+                decode_sampled_token_ids.append(
+                    sampler_output.sampled_token_ids.flatten())
+                decode_sampled_requests.extend(
+                    self.input_batch.req_ids[:num_decodes])
             htorch.core.mark_step()
+            if self.is_driver_worker and self.profiler.enabled:
+                # Stop recording 'execute_model' event
+                self.profiler.end()
+                event_end = self.profiler.get_timestamp_us()
+                counters = self.profiler_counter_helper.get_counter_dict(
+                    cache_config=self.cache_config,
+                    duration=event_end - self.event_start,
+                    seq_len=self._seq_len(decode_data.attn_metadata),
+                    batch_size_padded= \
+                        decode_data.token_ids.size(0),  # type: ignore
+                    real_batch_size=decode_data.num_decodes,
+                    prompt_batch_idx=None,
+                    is_prompt=False)
+                self.profiler.record_counter(self.event_start, counters)
         # From this point onward, all operations are done on CPU.
         # We already have tokens. Let's copy the data to
         # CPU as is, and then discard padded tokens.
-
-        prefill_sampled_token_ids = [
-            tensor.cpu() for tensor in prefill_sampled_token_ids
-        ]
-        decode_sampled_token_ids = [
-            tensor.cpu()[:num_decodes] for tensor in decode_sampled_token_ids
-        ]
-        sampled_token_ids_list = torch.cat(decode_sampled_token_ids +
-                                           prefill_sampled_token_ids).tolist()
-        sampled_token_requests = \
-            decode_sampled_requests + prefill_sampled_requests
-        max_req_index = max(self.input_batch.req_id_to_index.values())
-        postprocessed_sampled_token_ids: list[list]
-        postprocessed_sampled_token_ids = [[]
-                                           for _ in range(max_req_index + 1)]
-        for tok_id, req_id in zip(sampled_token_ids_list,
-                                  sampled_token_requests):
-            postprocessed_sampled_token_ids[
-                self.input_batch.req_id_to_index[req_id]].append(tok_id)
+        with self.profiler.record_event('internal', "sampler_postprocessing"):
+            prefill_sampled_token_ids = [
+                tensor.cpu() for tensor in prefill_sampled_token_ids
+            ]
+            decode_sampled_token_ids = [
+                tensor.cpu()[:num_decodes]
+                for tensor in decode_sampled_token_ids
+            ]
+            sampled_token_ids_list = torch.cat(
+                decode_sampled_token_ids + prefill_sampled_token_ids).tolist()
+            sampled_token_requests = \
+                decode_sampled_requests + prefill_sampled_requests
+            max_req_index = max(self.input_batch.req_id_to_index.values())
+            postprocessed_sampled_token_ids: list[list]
+            postprocessed_sampled_token_ids = [[]
+                                               for _ in range(max_req_index +
+                                                              1)]
+            for tok_id, req_id in zip(sampled_token_ids_list,
+                                      sampled_token_requests):
+                postprocessed_sampled_token_ids[
+                    self.input_batch.req_id_to_index[req_id]].append(tok_id)
 
         # NOTE(kzawora): idk what happens if part of batch doesn't have logprobs
 
@@ -1796,6 +1865,14 @@ def warmup_scenario(self,
         slot_mapping_device = _async_h2d_tensor_copy(slot_mapping, self.device)
 
         use_graphs = self._use_graphs()
+        phase = "prompt" if is_prompt else "decode"
+        scenario_name = ("warmup_"
+                         f"{phase}_"
+                         f"bs{batch_size}_"
+                         f"seq{query_seq_len}_"
+                         f"ctx{num_blocks}_"
+                         f"graphs{'T' if use_graphs else 'F'}")
+
         input_ids = torch.zeros((batch_size, query_seq_len),
                                 dtype=torch.int32,
                                 device='cpu')
@@ -1809,6 +1886,7 @@ def warmup_scenario(self,
         input_ids_device = _async_h2d_tensor_copy(input_ids, self.device)
         position_ids_device = _async_h2d_tensor_copy(position_ids, self.device)
         slot_mapping_device = _async_h2d_tensor_copy(slot_mapping, self.device)
+        self.profiler.start('internal', scenario_name)
         times = 3 if use_graphs or is_pt_profiler_run else 1
         for time_index in range(times):
             if is_prompt:
@@ -1882,6 +1960,7 @@ def warmup_scenario(self,
         }  # NOTE(kzawora): idk what to set here
         max_num_logprobs = 0  # NOTE(kzawora): idk what to set here
         # NOTE(kzawora: do this in a smarter way)
+        self.profiler.end()
         return None
         htorch.core.mark_step()
         sampling_metadata = SamplingMetadata(
@@ -2138,6 +2217,7 @@ def warmup_model(self) -> None:
             logger.info("Skipping warmup...")
             return
 
+        self.profiler.start('internal', 'warmup')
         start_mem = HabanaMemoryProfiler.current_device_memory_usage()
         start_time = time.perf_counter()
 
@@ -2189,6 +2269,7 @@ def warmup_model(self) -> None:
             f"Warmup finished in {elapsed_time:.0f} secs, "
             f"allocated {format_bytes(end_mem - start_mem)} of device memory")
         logger.info(msg)
+        self.profiler.end()
 
     def shutdown_inc(self):
         can_finalize_inc = self._is_quant_with_inc() and \