add speculative metrics

Author: kris1025

tensorrt_llm/bench/dataclasses/reporting.py

@@ -72,7 +72,7 @@ def register_request_perf_item(self, request_perf_item: PerfItemTuple):
         if request_perf_item.response_is_final:
             self.num_complete = self.num_complete + 1
 
-    def generate_statistics_summary(self) -> None:
+    def generate_statistics_summary(self, num_draft_tokens: int) -> None:
         """Generate summary statistics from internally stored statistics.
 
         Returns:
@@ -90,41 +90,59 @@ def generate_statistics_summary(self) -> None:
 
         intertoken_avg_latencies = []
         output_tokens = []
-        request_acceptance = []
         total_decoding_iterations = 0
         ttft_times = []
         last_queue_time = 0.0
         queue_time_total = 0.0
 
+        num_draft_tokens = []
+        num_accepted_drafts = []
+        draft_acceptance_rate = []
+        acceptance_length = []
+
         for entry in self.requests.values():
             start_time = min(entry.start_timestamp, start_time)
             end_time = max(entry.end_timestamp, end_time)
             last_queue_time = max(entry.start_timestamp, last_queue_time)
-            request_ar = entry.num_generated_tokens / (entry.decode_iteration +
-                                                       1)
 
             request_latencies.append(entry.end_to_end_latency)
             generation_latencies.append(entry.generation_time)
             generation_throughputs.append(entry.generation_token_throughput)
             ttft_times.append(entry.time_to_first_token)
             intertoken_avg_latencies.append(entry.intertoken_latency)
-            request_acceptance.append(request_ar)
+            
             output_throughput_per_user.append(entry.output_token_throughput)
             total_decoding_iterations += entry.decode_iteration + 1
 
             output_tokens.append(entry.num_total_output_tokens)
             total_input_tokens += entry.num_input_tokens
 
-        global_acceptance_rate = sum(output_tokens) / total_decoding_iterations
+            # For speculative decoding, we need to track the number of draft tokens per request and the number of accepted draft tokens per request
+            if num_draft_tokens > 0:
+                num_draft_tokens.append(num_draft_tokens * (entry.decode_iteration + 1))
+                num_accepted_draft_tokens.append(entry.num_generated_tokens - entry.decode_iteration - 1)
+                draft_acceptance_rate.append(float(num_accepted_draft_tokens[-1]) / float(num_draft_tokens[-1]))
+                acceptance_length.append(entry.num_generated_tokens / (entry.decode_iteration +
+                                                       1))
+
+        global_acceptance_length = sum(acceptance_length) / len(acceptance_length)
         queue_time_total = last_queue_time - start_time
-        percentile_request_accept = PercentileStats.from_iterable(
-            request_acceptance) if request_acceptance else None
+        
+        num_draft_tokens_percentiles = PercentileStats.from_iterable(
+            num_draft_tokens) if num_draft_tokens else None
+        num_accepted_drafts_percentiles = PercentileStats.from_iterable(
+            num_accepted_drafts) if num_accepted_drafts else None
+        draft_acceptance_rate_percentiles = PercentileStats.from_iterable(
+            draft_acceptance_rate) if draft_acceptance_rate else None
+        acceptance_length_percentiles = PercentileStats.from_iterable(
+            acceptance_length) if acceptance_length else None
 
         stats = BenchmarkStatistics(
             num_requests=num_requests,
             total_latency_ns=end_time - start_time,
             total_output_tokens=sum(output_tokens),
             total_input_tokens=total_input_tokens,
+            total_draft_tokens=sum(draft_tokens),
             acceptance_rate=global_acceptance_rate,
             request_latency_percentiles=PercentileStats.from_iterable(
                 request_latencies),
@@ -139,7 +157,11 @@ def generate_statistics_summary(self) -> None:
                 generation_latencies),
             token_percentiles=PercentileStats.from_iterable(output_tokens),
             issue_rate_ns=queue_time_total / num_requests,
-            acceptance_percentiles=percentile_request_accept,
+            acceptance_length=global_acceptance_length,
+            num_draft_tokens_percentiles=num_draft_tokens_percentiles,
+            num_accepted_drafts_percentiles=num_accepted_drafts_percentiles,
+            draft_acceptance_rate_percentiles=draft_acceptance_rate_percentiles,
+            acceptance_length_percentiles=acceptance_length_percentiles,
         )
 
         return stats
@@ -171,6 +193,7 @@ def __init__(self,
         self.logger = logger
         self.kwargs = kwargs
         self.streaming = streaming
+        
 
     @staticmethod
     def convert_to_ms(ns: float) -> float:
@@ -321,7 +344,9 @@ def get_statistics_dict(self) -> Dict[str, Any]:
             "avg_num_concurrent_requests":
             self.statistics.avg_concurrent_requests,
             "avg_input_length": self.statistics.average_input_length,
-            "avg_output_length": self.statistics.average_output_length
+            "avg_output_length": self.statistics.average_output_length,
+            "total_draft_tokens": self.statistics.total_draft_tokens,
+            "avg_draft_tokens_per_request": 0
         }
 
         # Performance stats
@@ -415,9 +440,18 @@ def get_statistics_dict(self) -> Dict[str, Any]:
             stats_dict["decoding_stats"] = {
                 "mode":
                 decoding_mode,
-                "acceptance_percentiles":
-                self.statistics.acceptance_percentiles.model_dump(
-                    exclude_none=True, by_alias=True, mode='json')
+                "draft_tokens_percentiles":
+                self.statistics.num_draft_tokens_percentiles.model_dump(
+                    exclude_none=True, by_alias=True, mode='json') if self.statistics.num_draft_tokens_percentiles else None,
+                "num_accepted_drafts_percentiles":
+                self.statistics.num_accepted_drafts_percentiles.model_dump(
+                    exclude_none=True, by_alias=True, mode='json') if self.statistics.num_accepted_drafts_percentiles else None,
+                "draft_acceptance_rate_percentiles":
+                self.statistics.draft_acceptance_rate_percentiles.model_dump(
+                    exclude_none=True, by_alias=True, mode='json') if self.statistics.draft_acceptance_rate_percentiles else None,
+                "acceptance_length_percentiles":
+                self.statistics.acceptance_length_percentiles.model_dump(
+                    exclude_none=True, by_alias=True, mode='json') if self.statistics.acceptance_length_percentiles else None
             }
         # Dataset metadata
         stats_dict["dataset"] = self.dataset_metadata.model_dump(by_alias=True,
@@ -503,6 +537,8 @@ def report_statistics(self) -> None:
             f"Number of concurrent requests:  {requests['avg_num_concurrent_requests']:.4f}\n"
             f"Average Input Length (tokens):  {requests['avg_input_length']:.4f}\n"
             f"Average Output Length (tokens): {requests['avg_output_length']:.4f}\n"
+            f"Total Draft Tokens:             {requests['total_draft_tokens']}\n"
+            f"Max Draft Tokens per Request:   {requests['avg_draft_tokens_per_request']}\n"
         )
 
         perf_header = (
@@ -557,19 +593,44 @@ def report_statistics(self) -> None:
         decoding_stats = ""
         if decoding is not None:
             decoding = stats_dict["decoding_stats"]
-            acc = decoding["acceptance_percentiles"]
-            acc_stats = "\n".join(
-                f"[AR] {key.upper():<7}: {acc[key]:.2f}" for key in
+            
+            num_draft_tokens = decoding["num_draft_tokens_percentiles"]
+            num_draft_tokens_stats = "\n".join(
+                f"[DT] {key.upper():<7}: {num_draft_tokens[key]:.2f}" for key in
+                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])
+
+            num_accepted_draft_tokens = decoding["num_accepted_drafts_percentiles"]
+            num_accepted_draft_tokens_stats = "\n".join(
+                f"[ADT] {key.upper():<7}: {num_accepted_draft_tokens[key]:.2f}" for key in
+                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])
+
+            draft_acceptance_rate = decoding["draft_acceptance_rate_percentiles"]
+            draft_acceptance_rate_stats = "\n".join(
+                f"[DAR] {key.upper():<7}: {draft_acceptance_rate[key]:.2f}" for key in
+                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])
+
+            acceptance_length = decoding["acceptance_length_percentiles"]
+            acceptance_length_stats = "\n".join(
+                f"[AL] {key.upper():<7}: {acceptance_length[key]:.2f}" for key in
                 ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])
 
             decoding_stats = (
                 "===========================================================\n"
                 f"= DECODING STATISTICS ({decoding['mode']})\n"
                 "===========================================================\n"
                 "\n"
-                "-- Acceptance Rate Details --------------------------------\n\n"
+                "-- Number of Draft Tokens Details --------------------------------\n\n"
                 "\n"
-                f"{acc_stats}"
+                f"{num_draft_tokens_stats}"
+                f"\n"
+                "-- Number of Accepted Draft Tokens Details --------------------------------\n\n"
+                f"{num_accepted_draft_tokens_stats}"
+                f"\n"
+                "-- Draft Acceptance Rate Details --------------------------------\n\n"
+                f"{draft_acceptance_rate_stats}"
+                f"\n"
+                "-- Acceptance Length Details --------------------------------\n\n"
+                f"{acceptance_length_stats}"
                 f"\n"
                 "===========================================================\n")
 
@@ -582,3 +643,29 @@ def report_statistics(self) -> None:
                         f"{self.dataset_metadata.get_summary_for_print()}")
         self.logger.info(logging_info)
         return self.statistics
+
+    def get_max_draft_len(self) -> int:
+        """Get max_draft_len from rt_cfg or speculative_config."""
+        # Try to get from rt_cfg first
+        if (self.rt_cfg.decoding_config
+                and self.rt_cfg.decoding_config.decoding_mode
+                != SpeculativeDecodingMode.NONE):
+            # For C++ backend, max_draft_len is stored in the engine config
+            # We need to read it from the engine config file
+            if self.rt_cfg.engine_dir:
+                config_path = self.rt_cfg.engine_dir / "config.json"
+                try:
+                    with open(config_path, "r") as config:
+                        engine_config = json.load(config)
+                    build_cfg = engine_config["build_config"]
+                    return build_cfg.get("max_draft_len", 0)
+                except (FileNotFoundError, KeyError, json.JSONDecodeError):
+                    pass
+        
+        # Try to get from speculative_config
+        if ("speculative_config" in self.kwargs
+                and self.kwargs["speculative_config"] is not None):
+            return self.kwargs["speculative_config"].max_draft_len or 0
+        
+        return 0
+

tensorrt_llm/bench/dataclasses/statistics.py

@@ -121,13 +121,14 @@ class BenchmarkStatistics(BaseModel):
     # Token-related Properties
     total_output_tokens: int
     total_input_tokens: int
+    total_draft_tokens: int = 0
 
     # General Information
     num_requests: int
     issue_rate_ns: float
 
     # Speculative Information
-    acceptance_rate: float
+    acceptance_length: float
 
     # Percentile-related Statistics
     request_latency_percentiles: Optional[PercentileStats] = None
@@ -137,7 +138,11 @@ class BenchmarkStatistics(BaseModel):
     ttft_percentiles: Optional[PercentileStats] = None
     generation_tp_percentiles: Optional[PercentileStats] = None
     generation_latency_percentiles: Optional[PercentileStats] = None
-    acceptance_percentiles: Optional[PercentileStats] = None
+    # Percentile-related Speculative Statistics
+    num_draft_tokens_percentiles: Optional[PercentileStats] = None
+    num_accepted_drafts_percentiles: Optional[PercentileStats] = None
+    draft_acceptance_rate_percentiles: Optional[PercentileStats] = None
+    acceptance_length_percentiles: Optional[PercentileStats] = None
 
     @computed_field
     def sum_per_request_latencies_ns(self) -> float: