removed v1

Signed-off-by: Settheworldonfireiii <[email protected]>

Author: Settheworldonfireiii

vllm/engine/llm_engine.py

@@ -1102,10 +1102,22 @@ def _process_model_outputs(self,
                 continue
 
             output: List[SequenceGroupOutput]
+            return_hidden_states = False
             if has_multiple_outputs:
                 output = outputs_by_sequence_group[i]
+                if self.model_config.task == "generate" and hasattr(
+                        outputs_by_sequence_group[0][0], 'hidden_states'):
+                    return_hidden_states = True
+                    for k in range(len(outputs_by_sequence_group[i])):
+                        output[k].hidden_states = outputs_by_sequence_group[i][
+                            k].hidden_states
             else:
                 output = [outputs_by_sequence_group[0][i]]
+                if self.model_config.task == "generate" and hasattr(
+                        outputs_by_sequence_group[0], 'hidden_states'):
+                    return_hidden_states = True
+                    output[0].hidden_states = outputs_by_sequence_group[
+                        0].hidden_states
 
             if not is_async:
                 if self.scheduler_config.is_multi_step:
@@ -1152,10 +1164,17 @@ def _process_model_outputs(self,
             seq_group.maybe_set_first_token_time(now)
             if not seq_group.is_prefill():
                 seq_group.set_last_token_time(now)
-            request_output = RequestOutputFactory.create(
-                seq_group,
-                self.seq_id_to_seq_group,
-                use_cache=self.use_cached_outputs)
+            if return_hidden_states:
+                request_output = RequestOutputFactory.create(
+                    seq_group,
+                    self.seq_id_to_seq_group,
+                    use_cache=self.use_cached_outputs,
+                    hidden_states=output[0].hidden_states)
+            else:
+                request_output = RequestOutputFactory.create(
+                    seq_group,
+                    self.seq_id_to_seq_group,
+                    use_cache=self.use_cached_outputs)
             if request_output:
                 ctx.request_outputs.append(request_output)
 

vllm/outputs.py

@@ -118,6 +118,7 @@ def __init__(
         encoder_prompt: Optional[str] = None,
         encoder_prompt_token_ids: Optional[list[int]] = None,
         num_cached_tokens: Optional[int] = None,
+        hidden_states: Optional[torch.Tensor] = None,
         *,
         multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,
     ) -> None:
@@ -133,6 +134,8 @@ def __init__(
         self.encoder_prompt = encoder_prompt
         self.encoder_prompt_token_ids = encoder_prompt_token_ids
         self.num_cached_tokens = num_cached_tokens
+        if hidden_states is not None:
+            self.hidden_states = hidden_states
 
     def add(self, next_output: "RequestOutput") -> None:
         """Merge subsequent RequestOutput into this one"""
@@ -160,8 +163,11 @@ def add(self, next_output: "RequestOutput") -> None:
 
     @classmethod
     def from_seq_group(
-        cls, seq_group: SequenceGroup, use_cache: bool,
-        seq_id_to_seq_group: dict[str, SequenceGroupBase]
+        cls,
+        seq_group: SequenceGroup,
+        use_cache: bool,
+        seq_id_to_seq_group: dict[str, SequenceGroupBase],
+        hidden_states: Optional[torch.Tensor] = None,
     ) -> Optional["RequestOutput"]:
         finished = seq_group.is_finished()
 
@@ -291,21 +297,37 @@ def from_seq_group(
             prompt_logprobs = None
         finished_time = time.time() if finished else None
         seq_group.set_finished_time(finished_time)
-
-        init_kwargs = {
-            "request_id": seq_group.request_id,
-            "prompt": prompt,
-            "prompt_token_ids": prompt_token_ids,
-            "prompt_logprobs": prompt_logprobs,
-            "outputs": outputs,
-            "finished": finished,
-            "metrics": seq_group.metrics,
-            "lora_request": seq_group.lora_request,
-            "encoder_prompt": encoder_prompt,
-            "encoder_prompt_token_ids": encoder_prompt_token_ids,
-            "num_cached_tokens": num_cached_tokens,
-            "multi_modal_placeholders": seq_group.multi_modal_placeholders
-        }
+        if hidden_states is not None:
+            init_kwargs = {
+                "request_id": seq_group.request_id,
+                "prompt": prompt,
+                "prompt_token_ids": prompt_token_ids,
+                "prompt_logprobs": prompt_logprobs,
+                "outputs": outputs,
+                "finished": finished,
+                "metrics": seq_group.metrics,
+                "lora_request": seq_group.lora_request,
+                "encoder_prompt": encoder_prompt,
+                "encoder_prompt_token_ids": encoder_prompt_token_ids,
+                "num_cached_tokens": num_cached_tokens,
+                "multi_modal_placeholders": seq_group.multi_modal_placeholders,
+                "hidden_states": hidden_states,
+            }
+        else:
+            init_kwargs = {
+                "request_id": seq_group.request_id,
+                "prompt": prompt,
+                "prompt_token_ids": prompt_token_ids,
+                "prompt_logprobs": prompt_logprobs,
+                "outputs": outputs,
+                "finished": finished,
+                "metrics": seq_group.metrics,
+                "lora_request": seq_group.lora_request,
+                "encoder_prompt": encoder_prompt,
+                "encoder_prompt_token_ids": encoder_prompt_token_ids,
+                "num_cached_tokens": num_cached_tokens,
+                "multi_modal_placeholders": seq_group.multi_modal_placeholders,
+            }
 
         if use_cache:
             request_output = seq_group.cached_request_output
@@ -385,12 +407,18 @@ class RequestOutputFactory:
     @staticmethod
     def create(seq_group: SequenceGroup,
                seq_id_to_seq_group: dict[str, SequenceGroupBase],
-               use_cache: bool = False):
+               use_cache: bool = False,
+               hidden_states: Optional[torch.Tensor] = None):
         if seq_group.pooled_data is not None:
             return PoolingRequestOutput.from_seq_group(seq_group)
         else:
-            return RequestOutput.from_seq_group(seq_group, use_cache,
-                                                seq_id_to_seq_group)
+            if hidden_states is not None:
+                return RequestOutput.from_seq_group(seq_group, use_cache,
+                                                    seq_id_to_seq_group,
+                                                    hidden_states)
+            else:
+                return RequestOutput.from_seq_group(seq_group, use_cache,
+                                                    seq_id_to_seq_group)
 
 
 @dataclass

vllm/sampling_params.py

@@ -186,6 +186,8 @@ class SamplingParams(
         allowed_token_ids: If provided, the engine will construct a logits
             processor which only retains scores for the given token ids.
             Defaults to None.
+        return_hidden_states: If provided, hidden states of the last attention
+            block are returned in the output
         extra_args: Arbitrary additional args, that can be used by custom
             sampling implementations. Not used by any in-tree sampling
             implementations.
@@ -233,6 +235,9 @@ class SamplingParams(
     allowed_token_ids: Optional[list[int]] = None
     extra_args: Optional[dict[str, Any]] = None
 
+    # Output hidden states or not
+    return_hidden_states: Optional[bool] = None
+
     # Fields used for bad words
     bad_words: Optional[list[str]] = None
     _bad_words_token_ids: Optional[list[list[int]]] = None

vllm/sequence.py

@@ -1093,6 +1093,7 @@ class CompletionSequenceGroupOutput(
     # Prompt logprob for each prompt query token.
     prompt_logprobs: Optional[PromptLogprobs]
     step_index: Optional[int] = 0
+    hidden_states: Optional[torch.Tensor] = None
 
     def __repr__(self) -> str:
         return (f"CompletionSequenceGroupOutput(samples={self.samples}, "

vllm/worker/model_runner.py

@@ -1714,6 +1714,19 @@ def execute_model(
         # virtual engines share the same kv cache.
         virtual_engine = model_input.virtual_engine
         previous_hidden_states = kwargs.get("previous_hidden_states")
+
+        # overrides self.return_hidden_states that was
+        # assigned during initialization
+        # the rationale is giving users the option
+        # to receive hidden states or not
+        # from the same model w/o re-init it
+        if (model_input.sampling_metadata is not None
+                and hasattr(model_input.sampling_metadata, 'seq_groups')
+                and model_input.sampling_metadata.seq_groups is not None):
+            self.return_hidden_states = (
+                model_input.sampling_metadata.seq_groups[0].sampling_params.
+                return_hidden_states)
+
         if prefill_meta is None and decode_meta.use_cuda_graph:
             assert model_input.input_tokens is not None
             graph_batch_size = model_input.input_tokens.shape[0]