add request level, per-step acceptance counts tracking for spec dec

Signed-off-by: Bryan Lu <[email protected]>

Author: luyuzhe111

examples/offline_inference/eagle.py

@@ -7,89 +7,113 @@
 
 from vllm import LLM, SamplingParams
 
-parser = argparse.ArgumentParser()
-
-parser.add_argument(
-    "--dataset",
-    type=str,
-    default="./examples/data/gsm8k.jsonl",
-    help="downloaded from the eagle repo " \
-    "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/"
-)
-parser.add_argument("--max_num_seqs", type=int, default=8)
-parser.add_argument("--num_prompts", type=int, default=80)
-parser.add_argument("--num_spec_tokens", type=int, default=2)
-parser.add_argument("--tp", type=int, default=1)
-parser.add_argument("--draft_tp", type=int, default=1)
-parser.add_argument("--enforce_eager", action='store_true')
-parser.add_argument("--enable_chunked_prefill", action='store_true')
-parser.add_argument("--max_num_batched_tokens", type=int, default=2048)
-parser.add_argument("--temp", type=float, default=0)
-
-args = parser.parse_args()
-
-print(args)
-
-model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
-eagle_dir = "abhigoyal/EAGLE-LLaMA3-Instruct-8B-vllm"
-
-max_model_len = 2048
-
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
-if os.path.exists(args.dataset):
-    prompts = []
-    num_prompts = args.num_prompts
-    with open(args.dataset) as f:
-        for line in f:
-            data = json.loads(line)
-            prompts.append(data["turns"][0])
-else:
-    prompts = ["The future of AI is", "The president of the United States is"]
-
-prompts = prompts[:args.num_prompts]
-num_prompts = len(prompts)
-
-prompt_ids = [
-    tokenizer.apply_chat_template([{
-        "role": "user",
-        "content": prompt
-    }],
-                                  add_generation_prompt=True)
-    for prompt in prompts
-]
-
-llm = LLM(
-    model=model_dir,
-    trust_remote_code=True,
-    tensor_parallel_size=args.tp,
-    enable_chunked_prefill=args.enable_chunked_prefill,
-    max_num_batched_tokens=args.max_num_batched_tokens,
-    enforce_eager=args.enforce_eager,
-    max_model_len=max_model_len,
-    max_num_seqs=args.max_num_seqs,
-    gpu_memory_utilization=0.8,
-    speculative_config={
-        "model": eagle_dir,
-        "num_speculative_tokens": args.num_spec_tokens,
-        "draft_tensor_parallel_size": args.draft_tp,
-        "max_model_len": max_model_len,
-    },
-    disable_log_stats=False,
-)
-
-sampling_params = SamplingParams(temperature=args.temp, max_tokens=256)
-
-outputs = llm.generate(prompt_token_ids=prompt_ids,
-                       sampling_params=sampling_params)
-
-# calculate the average number of accepted tokens per forward pass, +1 is
-# to account for the token from the target model that's always going to be
-# accepted
-acceptance_counts = [0] * (args.num_spec_tokens + 1)
-for output in outputs:
-    for step, count in enumerate(output.metrics.spec_token_acceptance_counts):
-        acceptance_counts[step] += count
-
-print(f"mean acceptance length: \
-    {sum(acceptance_counts) / acceptance_counts[0]:.2f}")
+
+def load_prompts(dataset_path, num_prompts):
+    if os.path.exists(dataset_path):
+        prompts = []
+        try:
+            with open(dataset_path) as f:
+                for line in f:
+                    data = json.loads(line)
+                    prompts.append(data["turns"][0])
+        except Exception as e:
+            print(f"Error reading dataset: {e}")
+            return []
+    else:
+        prompts = [
+            "The future of AI is", "The president of the United States is"
+        ]
+
+    return prompts[:num_prompts]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="./examples/data/gsm8k.jsonl",
+        help="downloaded from the eagle repo " \
+        "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/"
+    )
+    parser.add_argument("--max_num_seqs", type=int, default=8)
+    parser.add_argument("--num_prompts", type=int, default=80)
+    parser.add_argument("--num_spec_tokens", type=int, default=2)
+    parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--draft_tp", type=int, default=1)
+    parser.add_argument("--enforce_eager", action='store_true')
+    parser.add_argument("--enable_chunked_prefill", action='store_true')
+    parser.add_argument("--max_num_batched_tokens", type=int, default=2048)
+    parser.add_argument("--temp", type=float, default=0)
+    parser.add_argument("--use_v1", type=str, default="1", help='1 or 0')
+    args = parser.parse_args()
+
+    # TODO: remove this option once EAGLE in v1 is ready.
+    os.environ["VLLM_USE_V1"] = args.use_v1
+
+    model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
+    eagle_dir = "abhigoyal/EAGLE-LLaMA3-Instruct-8B-vllm"
+
+    max_model_len = 2048
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+    prompts = load_prompts(args.dataset, args.num_prompts)
+
+    prompt_ids = [
+        tokenizer.apply_chat_template([{
+            "role": "user",
+            "content": prompt
+        }],
+                                      add_generation_prompt=True)
+        for prompt in prompts
+    ]
+
+    llm = LLM(
+        model=model_dir,
+        trust_remote_code=True,
+        tensor_parallel_size=args.tp,
+        enable_chunked_prefill=args.enable_chunked_prefill,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        enforce_eager=args.enforce_eager,
+        max_model_len=max_model_len,
+        max_num_seqs=args.max_num_seqs,
+        gpu_memory_utilization=0.8,
+        speculative_config={
+            "method": "eagle",
+            "model": eagle_dir,
+            "num_speculative_tokens": args.num_spec_tokens,
+            "draft_tensor_parallel_size": args.draft_tp,
+            "max_model_len": max_model_len,
+        },
+        disable_log_stats=False,
+    )
+
+    sampling_params = SamplingParams(temperature=args.temp, max_tokens=256)
+
+    outputs = llm.generate(prompt_token_ids=prompt_ids,
+                           sampling_params=sampling_params)
+
+    # calculate the average number of accepted tokens per forward pass, +1 is
+    # to account for the token from the target model that's always going to be
+    # accepted
+    acceptance_counts = [0] * (args.num_spec_tokens + 1)
+    if args.use_v1 == '1':
+        for output in outputs:
+            for step, count in enumerate(
+                    output.spec_token_acceptance_counts[0]):
+                acceptance_counts[step] += count
+    else:
+        for output in outputs:
+            for step, count in enumerate(
+                    output.metrics.spec_token_acceptance_counts):
+                acceptance_counts[step] += count
+
+    print("-" * 50)
+    print(f"mean acceptance length: \
+        {sum(acceptance_counts) / acceptance_counts[0]:.2f}")
+    print("-" * 50)
+
+
+if __name__ == "__main__":
+    main()

vllm/outputs.py

@@ -43,6 +43,7 @@ class CompletionOutput:
     finish_reason: Optional[str] = None
     stop_reason: Union[int, str, None] = None
     lora_request: Optional[LoRARequest] = None
+    spec_token_acceptance_counts: Optional[list[int]] = None
 
     def finished(self) -> bool:
         return self.finish_reason is not None
@@ -133,6 +134,9 @@ def __init__(
         self.encoder_prompt = encoder_prompt
         self.encoder_prompt_token_ids = encoder_prompt_token_ids
         self.num_cached_tokens = num_cached_tokens
+        self.spec_token_acceptance_counts = [
+            o.spec_token_acceptance_counts for o in outputs
+        ]
 
     def add(self, next_output: "RequestOutput") -> None:
         """Merge subsequent RequestOutput into this one"""

vllm/v1/core/sched/scheduler.py

@@ -590,6 +590,10 @@ def update_from_output(
                     num_draft_tokens=len(scheduled_spec_token_ids),
                     num_accepted_tokens=len(generated_token_ids) - 1)
 
+                for i in range(len(generated_token_ids)):
+                    if request.spec_token_acceptance_counts is not None:
+                        request.spec_token_acceptance_counts[i] += 1
+
             cached_encoder_input_ids = (
                 self.encoder_cache_manager.get_cached_input_ids(request))
             # OPTIMIZATION: Avoid list(set) if the set is empty.
@@ -651,7 +655,9 @@ 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(),
+                        spec_token_acceptance_counts=request.
+                        spec_token_acceptance_counts))
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors

vllm/v1/engine/__init__.py

@@ -59,6 +59,7 @@ class EngineCoreRequest(
     eos_token_id: Optional[int]
     arrival_time: float
     lora_request: Optional[LoRARequest]
+    spec_token_acceptance_counts: Optional[list[int]]
 
 
 class EngineCoreEventType(enum.IntEnum):
@@ -101,6 +102,7 @@ class EngineCoreOutput(
     finish_reason: Optional[FinishReason] = None
     stop_reason: Union[int, str, None] = None
     events: Optional[list[EngineCoreEvent]] = None
+    spec_token_acceptance_counts: Optional[list[int]] = None
 
     @property
     def finished(self) -> bool:

vllm/v1/engine/llm_engine.py

@@ -183,11 +183,20 @@ def add_request(
         priority: int = 0,
     ) -> None:
         # Process raw inputs into the request.
-        request = self.processor.process_inputs(request_id, prompt, params,
-                                                arrival_time, lora_request,
-                                                trace_headers,
-                                                prompt_adapter_request,
-                                                priority)
+        num_spec_tokens = 0
+        if self.vllm_config.speculative_config is not None:
+            num_spec_tokens = (
+                self.vllm_config.speculative_config.num_speculative_tokens)
+        request = self.processor.process_inputs(
+            request_id,
+            prompt,
+            params,
+            arrival_time,
+            lora_request,
+            trace_headers,
+            prompt_adapter_request,
+            priority,
+            num_spec_tokens=num_spec_tokens)
 
         n = params.n if isinstance(params, SamplingParams) else 1
 

vllm/v1/engine/output_processor.py

@@ -136,10 +136,9 @@ def from_new_request(
         )
 
     def make_request_output(
-        self,
-        new_token_ids: list[int],
-        finish_reason: Optional[FinishReason],
+        self, new_token_ids: list[int], finish_reason: Optional[FinishReason],
         stop_reason: Union[int, str, None],
+        spec_token_acceptance_counts: Optional[list[int]]
     ) -> Optional[RequestOutput]:
 
         finished = finish_reason is not None
@@ -150,7 +149,10 @@ def make_request_output(
             return None
 
         completion_output = self._new_completion_output(
-            new_token_ids, finish_reason, stop_reason)
+            new_token_ids,
+            finish_reason,
+            stop_reason,
+            spec_token_acceptance_counts=spec_token_acceptance_counts)
 
         request_id = self.request_id
         if self.parent_req is None:
@@ -186,10 +188,9 @@ def _new_request_output(
         )
 
     def _new_completion_output(
-        self,
-        token_ids: list[int],
-        finish_reason: Optional[FinishReason],
-        stop_reason: Union[int, str, None],
+            self, token_ids: list[int], finish_reason: Optional[FinishReason],
+            stop_reason: Union[int, str, None],
+            spec_token_acceptance_counts: Optional[list[int]]
     ) -> CompletionOutput:
 
         finished = finish_reason is not None
@@ -212,7 +213,8 @@ def _new_completion_output(
             logprobs=logprobs,
             cumulative_logprob=self.logprobs_processor.cumulative_logprob,
             finish_reason=str(finish_reason) if finished else None,
-            stop_reason=stop_reason if finished else None)
+            stop_reason=stop_reason if finished else None,
+            spec_token_acceptance_counts=spec_token_acceptance_counts)
 
 
 class OutputProcessor:
@@ -337,7 +339,11 @@ def process_outputs(
 
             # 4) Create and handle RequestOutput objects.
             if request_output := req_state.make_request_output(
-                    new_token_ids, finish_reason, stop_reason):
+                    new_token_ids,
+                    finish_reason,
+                    stop_reason,
+                    spec_token_acceptance_counts=engine_core_output.
+                    spec_token_acceptance_counts):
                 if req_state.queue is not None:
                     # AsyncLLM: put into queue for handling by generate().
                     req_state.queue.put(request_output)

vllm/v1/engine/processor.py

@@ -176,6 +176,7 @@ def process_inputs(
         trace_headers: Optional[Mapping[str, str]] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
         priority: int = 0,
+        num_spec_tokens: int = 0,
     ) -> EngineCoreRequest:
 
         # TODO(woosuk): Support pooling models.
@@ -278,7 +279,7 @@ def process_inputs(
             eos_token_id=eos_token_id,
             arrival_time=arrival_time,
             lora_request=lora_request,
-        )
+            spec_token_acceptance_counts=[0] * (num_spec_tokens + 1))
 
     def _validate_model_inputs(self,
                                inputs: ProcessorInputs,