Libint/add samplemetatensorcache3 (#1991)

Co-authored-by: Libin Tang <[email protected]>

Author: yeonsily

vllm/model_executor/layers/sampler.py

@@ -199,6 +199,7 @@ def __init__(self):
         self.should_modify_greedy_probs_inplace = False
         # Add HPU cache class variables
         self._prompt_tokens_hpu_cache: Optional[torch.Tensor] = None
+        self._output_tokens_hpu_cache: Optional[torch.Tensor] = None
         self._cached_seq_ids: Optional[set] = None
 
     def _init_sampling_tensors(
@@ -222,7 +223,7 @@ def _init_sampling_tensors(
          top_k_scalar, top_p_scalar, current_seq_ids) = \
             SamplingTensors.from_sampling_metadata(
              sampling_metadata, vocab_size, logits.device, logits.dtype, \
-             self._prompt_tokens_hpu_cache, self._cached_seq_ids)
+             self._prompt_tokens_hpu_cache, self._output_tokens_hpu_cache, self._cached_seq_ids)
 
         self._sampling_tensors = sampling_tensors
         self._do_penalties = do_penalties
@@ -237,6 +238,7 @@ def _init_sampling_tensors(
         # After tensors are created, update cache
         if self._cached_seq_ids != current_seq_ids:
             self._prompt_tokens_hpu_cache = None
+            self._output_tokens_hpu_cache = None
             self._cached_seq_ids = current_seq_ids
 
     def forward(

vllm/model_executor/sampling_metadata.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-import torch
+import torch,time
 
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams, SamplingType
@@ -423,6 +423,7 @@ def from_sampling_metadata(
         device: torch.device,
         dtype: torch.dtype,
         prompt_tokens_cache: torch.tensor,
+        output_tokens_cache: torch.tensor,
         past_seq_ids: set,
     ) -> tuple["SamplingTensors", bool, bool, bool, Optional[int],
                Optional[float], Optional[torch.tensor]]:
@@ -516,7 +517,7 @@ def from_sampling_metadata(
                         current_seq_ids.update(seq_ids)
             if current_seq_ids != past_seq_ids:
                 prompt_tokens_cache = None
-
+                output_tokens_cache = None
         top_k_scalar = top_ks[0] if do_top_p_top_k and all(
             k == top_ks[0] for k in top_ks) else None
         top_p_scalar = top_ps[0] if do_top_p_top_k and all(
@@ -536,6 +537,7 @@ def from_sampling_metadata(
             device,
             dtype,
             prompt_tokens_cache,
+            output_tokens_cache,
         )
         return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p,
                 top_k_scalar, top_p_scalar, current_seq_ids)
@@ -556,6 +558,7 @@ def from_lists(
         device: torch.device,
         dtype: torch.dtype,
         prompt_tokens_cache: torch.tensor,
+        output_tokens_cache: torch.tensor,
     ) -> "SamplingTensors":
         # Note that the performance will be very bad without
         # pinned memory.
@@ -568,6 +571,14 @@ def from_lists(
                     prompt_tokens_cache.device == device):
                     # Reuse cached prompt_tokens already on HPU
                     prompt_t = prompt_tokens_cache
+                    # Get the last element from each list
+                    last_elements = [out[-1] for out in output_tokens]
+                    lengths = [len(out)-1 for out in output_tokens]
+                    indices = torch.tensor(lengths, device=device)
+                    rows = torch.arange(output_tokens_cache.shape[0], device=device)
+                    # Convert to a PyTorch tensor with shape [4, 1]
+                    last_elements_t = torch.tensor(last_elements).unsqueeze(1).to(output_tokens_cache.device)
+                    output_t = output_tokens_cache.index_put_((rows, indices), last_elements_t)
                 else:
                     prompt_t = make_tensor_with_pad_align(
                         prompt_tokens,
@@ -577,14 +588,14 @@ def from_lists(
                         pin_memory=pin_memory,
                         max_len_align=1024,
                     )
-                output_t = make_tensor_with_pad_align(
-                    output_tokens,
-                    vocab_size,
-                    device="cpu",
-                    dtype=torch.int64,
-                    pin_memory=pin_memory,
-                    max_len_align=1024,
-                )
+                    output_t = make_tensor_with_pad_align(
+                        output_tokens,
+                        vocab_size,
+                        device="cpu",
+                        dtype=torch.int64,
+                        pin_memory=pin_memory,
+                        max_len_align=1024,
+                    )
             else:
                 prompt_t = make_tensor_with_pad(
                     prompt_tokens,
@@ -649,7 +660,7 @@ def from_lists(
         )
         # Because the memory is pinned, we can do non-blocking
         # transfer to device.
-
+        output_t=output_t.to(device=device, non_blocking=True) if output_t.device != device else output_t
         return cls(
             temperatures=temperatures_t.to(device=device, non_blocking=True),
             top_ps=top_ps_t.to(device=device, non_blocking=True),
@@ -662,5 +673,5 @@ def from_lists(
             repetition_penalties=repetition_penalties_t.to(device=device,
                                                            non_blocking=True),
             prompt_tokens=prompt_t.to(device=device, non_blocking=True) if prompt_t.device != device else prompt_t,
-            output_tokens=output_t.to(device=device, non_blocking=True),
+            output_tokens=output_t
         )

vllm/worker/hpu_model_runner.py

@@ -1888,12 +1888,15 @@ def _prepare_prompt(
         if image_index_tensor is not None:
             multi_modal_kwargs['image_index'] = image_index_tensor
 
-        use_mediapipe = os.getenv("VLLM_USE_MEDIA_PIPELINE", "false").lower() in ("1", "true", "yes")
+        use_mediapipe = os.getenv("VLLM_USE_MEDIA_PIPELINE",
+                                  "false").lower() in ("1", "true", "yes")
         if use_mediapipe:
             # With mediapipe path some tensors will already be on HPU, we only move to HPU if needed
             for key in multi_modal_kwargs.keys():
-                if hasattr(multi_modal_kwargs[key], "device") and multi_modal_kwargs[key].device != self.device:
-                    multi_modal_kwargs[key] = self.move_to_device(multi_modal_kwargs[key])
+                if hasattr(multi_modal_kwargs[key], "device"
+                           ) and multi_modal_kwargs[key].device != self.device:
+                    multi_modal_kwargs[key] = self.move_to_device(
+                        multi_modal_kwargs[key])
         else:
             multi_modal_kwargs = MultiModalKwargs.as_kwargs(multi_modal_kwargs,
                                                             device=self.device)
@@ -4030,6 +4033,7 @@ def try_revert_dummy_output_tokens():
                     if sampling_tensors.prompt_tokens.numel() > 0:
                         # Cache the prompt_tokens tensor that's already on HPU
                         self.model.sampler._prompt_tokens_hpu_cache = sampling_tensors.prompt_tokens
+                        self.model.sampler._output_tokens_hpu_cache = sampling_tensors.output_tokens
                 if use_delayed_sampling \
                    and model_input.async_callback is not None:
                     model_input.async_callback()