Merge branch 'main' of github.com:character-tech/vllm

amogkam · amogkam · commit 2c1c8e354729 · 2025-05-07T19:21:44.000Z
diff --git a/vllm/model_executor/models/transformers.py b/vllm/model_executor/models/transformers.py
@@ -369,6 +369,15 @@ def forward(
 
         return hidden_states
 
+    def compute_additional_head(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        if get_pp_group().is_last_rank and hasattr(self.model,
+                                                   "compute_additional_head"):
+            return self.model.compute_additional_head(hidden_states)
+        return None
+
     def load_weights(self, weights: Iterable[tuple[str,
                                                    torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
@@ -463,6 +472,14 @@ def compute_logits(
                                        sampling_metadata)
         return logits
 
+    def compute_additional_head(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        if hasattr(self.model, "compute_additional_head"):
+            return self.model.compute_additional_head(hidden_states)
+        return None
+
     def load_weights(self, weights: Iterable[tuple[str,
                                                    torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
diff --git a/vllm/outputs.py b/vllm/outputs.py
@@ -12,8 +12,9 @@
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import MultiModalPlaceholderDict
 from vllm.sampling_params import RequestOutputKind
-from vllm.sequence import (PromptLogprobs, RequestMetrics, SampleLogprobs,
-                           SequenceGroup, SequenceGroupBase, SequenceStatus)
+from vllm.sequence import (AdditionalHeads, PromptLogprobs, RequestMetrics,
+                           SampleLogprobs, SequenceGroup, SequenceGroupBase,
+                           SequenceStatus)
 
 
 @dataclass
@@ -28,6 +29,8 @@ class CompletionOutput:
             output text.
         logprobs: The log probabilities of the top probability words at each
             position if the logprobs are requested.
+        additional_heads: The additional head outputs of the generated output 
+            text.
         finish_reason: The reason why the sequence is finished.
         stop_reason: The stop string or token id that caused the completion
             to stop, None if the completion finished for some other reason
@@ -43,6 +46,7 @@ class CompletionOutput:
     finish_reason: Optional[str] = None
     stop_reason: Union[int, str, None] = None
     lora_request: Optional[LoRARequest] = None
+    additional_heads: Optional[AdditionalHeads] = None
 
     def finished(self) -> bool:
         return self.finish_reason is not None
@@ -53,6 +57,7 @@ def __repr__(self) -> str:
                 f"token_ids={self.token_ids}, "
                 f"cumulative_logprob={self.cumulative_logprob}, "
                 f"logprobs={self.logprobs}, "
+                f"additional_heads={self.additional_heads}, "
                 f"finish_reason={self.finish_reason}, "
                 f"stop_reason={self.stop_reason})")
 
diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py
@@ -248,6 +248,9 @@ class SamplingParams(
     bad_words: Optional[list[str]] = None
     _bad_words_token_ids: Optional[list[list[int]]] = None
 
+    # Fields used for additional heads (e.g. classifiers)
+    additional_heads: Optional[bool] = None
+
     @staticmethod
     def from_optional(
         n: Optional[int] = 1,
diff --git a/vllm/sequence.py b/vllm/sequence.py
@@ -53,6 +53,7 @@ class Logprob:
 PromptLogprobs = list[Optional[dict[int, Logprob]]]
 # {token_id -> logprob} for each sequence group.
 SampleLogprobs = list[dict[int, Logprob]]
+AdditionalHeads = list[list[float]]
 
 
 class SequenceStatus(enum.IntEnum):
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
@@ -647,6 +647,8 @@ def update_from_output(
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
+        new_additional_head_outputs = \
+            model_runner_output.additional_head_outputs
 
         new_running: list[Request] = []
         outputs: list[EngineCoreOutput] = []
@@ -665,6 +667,13 @@ def update_from_output(
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index]
+            if new_additional_head_outputs:
+                head_outputs_list = \
+                    new_additional_head_outputs.additional_head_outputs
+                additional_head_outputs_per_request = \
+                    head_outputs_list[req_index]
+            else:
+                additional_head_outputs_per_request = None
 
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id))
@@ -751,7 +760,10 @@ def update_from_output(
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         stop_reason=request.stop_reason,
-                        events=request.take_events()))
+                        events=request.take_events(),
+                        new_additional_head_outputs=
+                        additional_head_outputs_per_request,
+                    ))
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py
@@ -12,7 +12,7 @@
 from vllm.multimodal.inputs import PlaceholderRange
 from vllm.sampling_params import SamplingParams
 from vllm.v1.metrics.stats import SchedulerStats
-from vllm.v1.outputs import LogprobsLists, LogprobsTensors
+from vllm.v1.outputs import LogprobsLists, LogprobsTensors, AdditionalHeadOutputsPerRequest
 
 # These are possible values of RequestOutput.finish_reason,
 # so form part of the external API.
@@ -101,6 +101,8 @@ class EngineCoreOutput(
 
     new_logprobs: Optional[LogprobsLists] = None
     new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None
+    new_additional_head_outputs: Optional[
+        AdditionalHeadOutputsPerRequest] = None
 
     finish_reason: Optional[FinishReason] = None
     stop_reason: Union[int, str, None] = None
diff --git a/vllm/v1/engine/additional_heads.py b/vllm/v1/engine/additional_heads.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from vllm.logger import init_logger
+from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class AdditionalHeadsProcessor:
+    """Processor for additional head outputs from the model.
+    
+    This class handles storing and managing additional head outputs
+    for generated tokens, similar to how LogprobsProcessor handles logprobs.
+    """
+
+    # Additional head outputs for this request
+    additional_head_outputs: list[list[float]]
+
+    @classmethod
+    def from_new_request(
+        cls,
+        request: EngineCoreRequest,
+    ) -> "AdditionalHeadsProcessor":
+        """Create a new AdditionalHeadsProcessor for a request.
+        
+        Args:
+            request: The engine core request to process additional heads for.
+        """
+        return cls(additional_head_outputs=[], )
+
+    def update_from_output(self, output: EngineCoreOutput) -> None:
+        """Update with additional head outputs from EngineCore.
+        
+        Args:
+            output: The engine core output containing new additional 
+                head outputs.
+        """
+        if output.new_additional_head_outputs is not None:
+            self.additional_head_outputs.append(
+                output.new_additional_head_outputs.additional_head_outputs)
diff --git a/vllm/v1/engine/logprobs.py b/vllm/v1/engine/logprobs.py
@@ -195,4 +195,4 @@ def update_from_output(self, output: EngineCoreOutput) -> None:
         if output.new_logprobs is not None:
             self._update_sample_logprobs(output.new_logprobs)
         if output.new_prompt_logprobs_tensors is not None:
-            self._update_prompt_logprobs(output.new_prompt_logprobs_tensors)
+            self._update_prompt_logprobs(output.new_prompt_logprobs_tensors)
diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py
@@ -10,6 +10,7 @@
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
+from vllm.v1.engine.additional_heads import AdditionalHeadsProcessor
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
 from vllm.v1.engine.parallel_sampling import ParentRequest
@@ -81,6 +82,7 @@ def __init__(
         prompt: Optional[str],
         prompt_token_ids: list[int],
         logprobs_processor: LogprobsProcessor,
+        additional_heads_processor: AdditionalHeadsProcessor,
         detokenizer: IncrementalDetokenizer,
         max_tokens_param: Optional[int],
         arrival_time: float,
@@ -96,6 +98,7 @@ def __init__(
         self.prompt_token_ids = prompt_token_ids
         self.prompt_len = len(prompt_token_ids)
         self.logprobs_processor = logprobs_processor
+        self.additional_heads_processor = additional_heads_processor
         self.detokenizer = detokenizer
         self.max_tokens_param = max_tokens_param
         self.is_prefilling = True
@@ -130,6 +133,8 @@ def from_new_request(
                 tokenizer=tokenizer,
                 request=request,
             ),
+            additional_heads_processor=AdditionalHeadsProcessor.
+            from_new_request(request=request, ),
             detokenizer=IncrementalDetokenizer.from_new_request(
                 tokenizer=tokenizer,
                 request=request,
@@ -211,11 +216,18 @@ def _new_completion_output(
         if delta and logprobs:
             logprobs = logprobs[-len(token_ids):]
 
+        # Prepare additional heads, based on delta mode
+        additional_heads = (
+            self.additional_heads_processor.additional_head_outputs or None)
+        if delta and additional_heads:
+            additional_heads = additional_heads[-len(token_ids):]
+
         return CompletionOutput(
             index=self.request_index,
             text=text,
             token_ids=token_ids,
             logprobs=logprobs,
+            additional_heads=additional_heads,
             cumulative_logprob=self.logprobs_processor.cumulative_logprob,
             finish_reason=str(finish_reason) if finished else None,
             stop_reason=stop_reason if finished else None)
@@ -345,8 +357,11 @@ def process_outputs(
                 finish_reason = FinishReason.STOP
                 stop_reason = stop_string
 
-            # 3) Compute sample and prompt logprobs for request, if required.
+            # 3) Compute sample and prompt logprobs as well as additional heads
+            # for request, if required.
             req_state.logprobs_processor.update_from_output(engine_core_output)
+            req_state.additional_heads_processor.update_from_output(
+                engine_core_output)
 
             # 4) Create and handle RequestOutput objects.
             if request_output := req_state.make_request_output(
diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py
@@ -6,6 +6,16 @@
 import torch
 
 
+class AdditionalHeadOutputsPerRequest(NamedTuple):
+    # num_additional_head_outputs
+    additional_head_outputs: list[float]
+
+
+class AdditionalHeadOutputs(NamedTuple):
+    # num_generated_tokens x num_additional_head_outputs
+    additional_head_outputs: list[Optional[AdditionalHeadOutputsPerRequest]]
+
+
 class LogprobsLists(NamedTuple):
 
     # [num_reqs, max_num_logprobs + 1]
@@ -100,6 +110,9 @@ class ModelRunnerOutput:
     # [prompt_len]
     prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]]
 
+    # num_reqs x num_generated_tokens x num_additional_head_outputs
+    additional_head_outputs: Optional[AdditionalHeadOutputs] = None
+
 
 EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
     req_ids=[],
diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py
@@ -203,6 +203,9 @@ def __init__(
         # that are currently in the prefill phase.
         self.num_prompt_logprobs: dict[str, int] = {}
 
+        # req_idx -> bool
+        self.run_additional_heads: dict[int, bool] = {}
+
         # To accumulate prompt logprobs tensor chunks across prefill steps.
         self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
 
@@ -314,6 +317,9 @@ def add_request(
         if sampling_params.logit_bias is not None:
             self.logit_bias[req_index] = sampling_params.logit_bias
 
+        if sampling_params.additional_heads:
+            self.run_additional_heads[req_index] = True
+
         if sampling_params.allowed_token_ids:
             self.has_allowed_token_ids.add(req_id)
             if self.allowed_token_ids_mask_cpu_tensor is None:
@@ -371,6 +377,7 @@ def remove_request(self, req_id: str) -> Optional[int]:
         self.generators.pop(req_index, None)
         self.num_logprobs.pop(req_id, None)
         self.num_prompt_logprobs.pop(req_id, None)
+        self.run_additional_heads.pop(req_index, None)
         self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
 
         # LoRA
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
@@ -34,7 +34,8 @@
 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
                                         KVCacheConfig, KVCacheSpec,
                                         SlidingWindowSpec)
-from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
+from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AdditionalHeadOutputs,
+                             AdditionalHeadOutputsPerRequest, LogprobsTensors,
                              ModelRunnerOutput)
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.rejection_sampler import RejectionSampler
@@ -1121,6 +1122,40 @@ def execute_model(
         sample_hidden_states = hidden_states[logits_indices]
         logits = self.model.compute_logits(sample_hidden_states, None)
 
+        additional_head_indices_mask = [
+            (i in self.input_batch.run_additional_heads)
+            for i in range(self.input_batch.num_reqs)
+        ]
+        run_additional_heads = any(additional_head_indices_mask)
+
+        if run_additional_heads:
+            assert hasattr(self.model, "compute_additional_head")
+
+            # NOTE: In theory not all logit indices need additional
+            # head outputs and we could save some flops by masking.
+            # In practice, this is a small number of flops and this
+            # is simpler/introduces less overhead.
+            additional_heads_tensor = self.model.compute_additional_head(
+                sample_hidden_states, )
+
+            # Should be num_decode_tokens x additional_head_size
+            assert len(additional_heads_tensor.shape) == 2
+
+            # Don't return the additional head outputs where they aren't needed.
+            additional_head_outputs = AdditionalHeadOutputs(
+                additional_head_outputs=[
+                    AdditionalHeadOutputsPerRequest(
+                        additional_head_outputs=
+                        additional_head_outputs_per_request, )
+                    if mask else None
+                    for additional_head_outputs_per_request, mask in zip(
+                        additional_heads_tensor.tolist(),
+                        additional_head_indices_mask)
+                ], )
+
+        else:
+            additional_head_outputs = None
+
         # Apply structured output bitmasks if present
         if scheduler_output.grammar_bitmask is not None:
             self.apply_grammar_bitmask(scheduler_output, logits)
@@ -1291,6 +1326,7 @@ def execute_model(
             spec_token_ids=spec_token_ids,
             logprobs=logprobs_lists,
             prompt_logprobs_dict=prompt_logprobs_dict,
+            additional_head_outputs=additional_head_outputs,
         )
 
     def generate_draft_token_ids(
