reporting.py

from __future__ import annotations

import json
import os
from collections import defaultdict
from typing import Any, Dict, List, NamedTuple, Optional

import torch

try:
    import pynvml
except ImportError:
    pynvml = None

from tensorrt_llm._torch.pyexecutor.model_loader import \
    validate_and_set_kv_cache_quant
from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig
from tensorrt_llm.bench.dataclasses.general import DatasetMetadata
from tensorrt_llm.bench.dataclasses.statistics import (BenchmarkStatistics,
                                                       PercentileStats,
                                                       RequestRecord)
from tensorrt_llm.llmapi import KvCacheConfig
from tensorrt_llm.logger import Logger
from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode


class PerfItemTuple(NamedTuple):
    """A tuple for recording new requests and their responses."""
    start_timestamp: int
    end_timestamp: int
    request_id: int
    num_input_tokens: int
    response_is_final: bool
    error: bool
    tokens: List[int]
    decoding_iteration: int
    time_on_first_token: int


class StatsKeeper:
    """A statistics keeper for benchmarking."""

    def __init__(self) -> None:
        self.requests: Dict[int, RequestRecord] = defaultdict(RequestRecord)
        self.num_complete: int = 0
        self.total_energy: Optional[float] = None

    def register_request(
        self,
        request_id: int,
        timestamp: int,
        num_tokens: int,
    ) -> None:
        """Register a new request.

        DEPRECATED after switching to LLM API.

        Args:
            request_id (int): Identifier of the request.
            timestamp (int): Timestamp of when the request was submitted.
            num_tokens (int): Number of input tokens in the request.
        """
        record = self.requests[request_id]
        record.num_input_tokens = num_tokens
        record.start_timestamp = timestamp

    def register_request_perf_item(self, request_perf_item: PerfItemTuple):
        """
        Register request perf items, used exclusively with LLM API.
        """
        record = self.requests[request_perf_item.request_id]
        record.id = request_perf_item.request_id
        record.num_input_tokens = request_perf_item.num_input_tokens
        record.start_timestamp = request_perf_item.start_timestamp
        record.register_event(request_perf_item.error,
                              request_perf_item.response_is_final,
                              request_perf_item.end_timestamp,
                              request_perf_item.decoding_iteration,
                              request_perf_item.tokens,
                              request_perf_item.time_on_first_token)
        if request_perf_item.response_is_final:
            self.num_complete = self.num_complete + 1

    def set_energy(self, energy: Optional[float]):
        """Set the total energy for the benchmark."""
        self.total_energy = energy

    def generate_statistics_summary(self, max_draft_tokens: int) -> None:
        """Generate summary statistics from internally stored statistics.

        Returns:
            BenchmarkStatistic: Benchmark run statistics.
        """
        total_input_tokens: int = 0
        num_requests = len(self.requests)
        start_time = float("inf")
        end_time = -1

        request_latencies = []
        generation_latencies = []
        generation_throughputs = []
        output_throughput_per_user = []

        intertoken_avg_latencies = []
        output_tokens = []
        total_decoding_iterations = 0
        ttft_times = []
        last_queue_time = 0.0
        queue_time_total = 0.0

        num_draft_tokens = []
        num_accepted_draft_tokens = []
        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_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)
            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

            # For speculative decoding, we need to track the number of draft tokens per request and the number of accepted draft tokens per request
            if max_draft_tokens > 0:
                num_draft_tokens.append(max_draft_tokens *
                                        (entry.decode_iteration + 1))
                num_accepted_draft_tokens.append(entry.num_total_output_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_total_output_tokens /
                                         (entry.decode_iteration + 1))

        global_acceptance_length = sum(
            output_tokens) / total_decoding_iterations
        queue_time_total = last_queue_time - start_time

        num_draft_tokens_percentiles = PercentileStats.from_iterable(
            num_draft_tokens) if num_draft_tokens else None
        num_accepted_draft_tokens_percentiles = PercentileStats.from_iterable(
            num_accepted_draft_tokens) if num_accepted_draft_tokens 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_energy=self.total_energy,
            request_latency_percentiles=PercentileStats.from_iterable(
                request_latencies),
            tpot_percentiles=PercentileStats.from_iterable(
                intertoken_avg_latencies),
            output_throughput_percentiles=PercentileStats.from_iterable(
                output_throughput_per_user),
            ttft_percentiles=PercentileStats.from_iterable(ttft_times),
            generation_tp_percentiles=PercentileStats.from_iterable(
                generation_throughputs),
            generation_latency_percentiles=PercentileStats.from_iterable(
                generation_latencies),
            token_percentiles=PercentileStats.from_iterable(output_tokens),
            issue_rate_ns=queue_time_total / num_requests,
            acceptance_length=global_acceptance_length,
            num_draft_tokens_percentiles=num_draft_tokens_percentiles,
            num_accepted_draft_tokens_percentiles=
            num_accepted_draft_tokens_percentiles,
            draft_acceptance_rate_percentiles=draft_acceptance_rate_percentiles,
            acceptance_length_percentiles=acceptance_length_percentiles,
        )

        return stats


class ReportUtility:
    """A utility for reporting statistics."""

    def __init__(self,
                 statistics: StatsKeeper,
                 dataset_metadata: DatasetMetadata,
                 rt_cfg: RuntimeConfig,
                 logger: Logger,
                 kwargs: Dict[str, Any],
                 streaming: bool = False) -> None:
        """Initialize the ReportingController.

        Args:
            statistics (StatsKeeper): A statistics container.
            dataset_metadata (DatasetMetadata): Metadata about the dataset.
            rt_cfg (RuntimeConfig): Configuration for the run.
            logger (Logger): A logger for logging.
            streaming (bool, optional): Streaming benchmark used. Defaults to False.
        """
        self.dataset_metadata = dataset_metadata
        self.rt_cfg = rt_cfg
        self.logger = logger
        self.kwargs = kwargs
        self.raw_statistics = statistics
        self.statistics = statistics.generate_statistics_summary(
            self.get_max_draft_len())
        self.streaming = streaming

    def _query_gpu_info(self) -> Dict[str, Any]:
        """Query first GPU info (all GPUs must be identical for TRT-LLM)."""
        if not torch.cuda.is_available():
            return None

        try:
            cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
            physical_idx = int(
                cuda_visible.split(",")[0].strip()) if cuda_visible else 0

            props = torch.cuda.get_device_properties(physical_idx)
            gpu_info = {
                "name":
                getattr(props, "name", "Unknown"),
                "memory.total":
                float(getattr(props, "total_memory", 0.0)) / (1024.0**3),
                "clocks.mem":
                None,
            }
            if pynvml:
                try:
                    # Memory clock information is not reported by torch, using NVML instead
                    pynvml.nvmlInit()
                    handle = pynvml.nvmlDeviceGetHandleByIndex(physical_idx)
                    clocks_mem = pynvml.nvmlDeviceGetMaxClockInfo(
                        handle, pynvml.NVML_CLOCK_MEM) / 1000.0
                    gpu_info["clocks.mem"] = clocks_mem
                except pynvml.NVMLError as e:
                    self.logger.info(
                        f"Error querying GPU clock info with NVML: {e}")
                    gpu_info["clocks.mem"] = None
        except Exception as e:
            # broad catch for any other errors, since this is a non-critical operation
            self.logger.warning(f"Error querying GPU info: {e}")
            return None
        return gpu_info

    @staticmethod
    def convert_to_ms(ns: float) -> float:
        """Convert nanoseconds to milliseconds."""
        return ns * 1.0e-6

    @staticmethod
    def convert_to_s(ns: float) -> float:
        """Convert nanoseconds to seconds."""
        return ns * 1.0e-9

    @staticmethod
    def convert_rate_to_s(rate: float) -> float:
        """Convert rate to seconds."""
        return rate * 1.0e9

    @property
    def request_throughput_req_s(self) -> float:
        """Request throughput in requests per second."""
        return self.convert_rate_to_s(self.statistics.request_throughput_ns)

    @property
    def output_throughput_tok_s(self) -> float:
        """Output throughput in tokens per second."""
        return self.convert_rate_to_s(self.statistics.output_throughput_tok_ns)

    @property
    def total_token_throughput_tok_s(self) -> float:
        """Total token throughput in tokens per second."""
        return self.convert_rate_to_s(
            self.statistics.total_token_throughput_tok_ns)

    @property
    def per_user_generation_token_throughput_s(self) -> float:
        """Output throughput per user in tokens per second."""
        return self.convert_rate_to_s(
            self.statistics.per_user_generation_token_throughput_ns)

    @property
    def per_user_output_throughput_tok_s(self) -> float:
        """Output throughput per user in tokens per second."""
        return self.convert_rate_to_s(
            self.statistics.output_throughput_tok_ns_per_user)

    def get_output_tokens(self, tokenizer) -> Dict[int, List[str]]:
        retval = {}
        for req_id, request in self.raw_statistics.requests.items():
            output_str = tokenizer.decode(request.tokens)
            retval[req_id] = output_str
        return dict(sorted(retval.items()))

    def get_request_info(self, tokenizer) -> Dict[int, List[str]]:
        requests = []
        for request in self.raw_statistics.requests.values():
            entry = request.model_dump()
            entry["output"] = tokenizer.decode(entry["tokens"])
            entry["output_tokens"] = len(entry["tokens"])
            entry.pop("tokens")
            requests.append(entry)
        return requests

    def get_statistics_dict(self) -> Dict[str, Any]:
        """Get statistics as a dictionary.

        Returns:
            Dict[str, Any]: Dictionary containing benchmark statistics.
        """
        stats_dict = {
            "engine": {
                "model": self.rt_cfg.model,
                "model_path": str(self.rt_cfg.model_path),
                "engine_dir": str(self.rt_cfg.engine_dir),
                "revision": self.rt_cfg.revision,
                "version": self.rt_cfg.sw_version,
            },
        }

        # Machine / GPU details - query only first GPU (all GPUs must be identical)
        stats_dict["machine"] = self._query_gpu_info()

        # Retrieve KV cache information.
        kv_cache_config = self.kwargs.get("kv_cache_config", KvCacheConfig())
        if isinstance(kv_cache_config, KvCacheConfig):
            kv_cache_dtype = kv_cache_config.dtype
            kv_cache_mem_percent = kv_cache_config.free_gpu_memory_fraction
        elif isinstance(kv_cache_config, dict):
            kv_cache_dtype = kv_cache_config.get("dtype", "auto")
            kv_cache_mem_percent = kv_cache_config.get(
                "free_gpu_memory_fraction")
        else:
            raise ValueError(
                f"Invalid kv_cache_config type: {type(kv_cache_config)}.")

        kv_cache_mem_percent = kv_cache_mem_percent \
            if kv_cache_mem_percent is not None else None

        # Engine/Backend details
        if self.rt_cfg.backend not in ('pytorch', '_autodeploy'):
            config_path = self.rt_cfg.engine_dir / "config.json"
            with open(config_path, "r") as config:
                engine_config = json.load(config)
            build_cfg = engine_config["build_config"]
            pretrain_cfg = engine_config["pretrained_config"]

            stats_dict["engine"] |= {
                "backend":
                "TRT",
                "dtype":
                pretrain_cfg["dtype"],
                "kv_cache_dtype":
                pretrain_cfg["quantization"]["kv_cache_quant_algo"],
                "quantization":
                pretrain_cfg["quantization"]["quant_algo"],
                "max_input_length":
                build_cfg["max_input_len"],
                "max_sequence_length":
                build_cfg["max_seq_len"]
            }
        else:
            from tensorrt_llm._torch.model_config import ModelConfig
            from tensorrt_llm._utils import torch_dtype_to_str

            model = self.rt_cfg.model_path or self.rt_cfg.model
            model_config = ModelConfig.from_pretrained(model,
                                                       trust_remote_code=True)

            validate_and_set_kv_cache_quant(model_config, kv_cache_dtype)

            stats_dict["engine"] |= {
                "backend":
                "Pytorch",
                "dtype":
                torch_dtype_to_str(model_config.torch_dtype
                                   or model_config.pretrained_config.
                                   get_text_config().torch_dtype),
                "kv_cache_dtype":
                model_config.quant_config.kv_cache_quant_algo,
                "quantization":
                model_config.quant_config.quant_algo
            }

        # World and runtime info
        stats_dict["world_info"] = {
            "tp_size": self.rt_cfg.mapping["tp_size"],
            "pp_size": self.rt_cfg.mapping["pp_size"],
            "ep_size": self.rt_cfg.mapping["moe_ep_size"],
            "world_size": self.rt_cfg.mapping["world_size"],
            "max_batch_size": self.rt_cfg.settings_config.max_batch_size,
            "max_num_tokens": self.rt_cfg.settings_config.max_num_tokens,
            "scheduling_policy": self.rt_cfg.settings_config.scheduler_policy,
            "kv_cache_percentage": kv_cache_mem_percent,
            "issue_rate": self.convert_rate_to_s(self.statistics.issue_rate_ns)
        }

        # Request details
        stats_dict["request_info"] = {
            "num_requests": self.statistics.num_requests,
            "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
        }

        # Performance stats
        stats_dict["performance"] = {
            # End-to-End Latency (last request end - 1st request start)
            "total_latency_ms":
            self.convert_to_ms(self.statistics.total_latency_ns),
            # Average per request latency (sum request latencies / num requests)
            "avg_request_latency_ms":
            self.convert_to_ms(
                self.statistics.request_latency_percentiles.average),
            # Request throughput (num requests / end-to-end latency)
            "request_throughput_req_s":
            self.request_throughput_req_s,
            # NOTE: All mention of "output" below is in reference to OSL tokens
            # including the first token.
            # Output throughput (total output (OSL) tokens / end-to-end latency)
            "system_output_throughput_tok_s":
            self.output_throughput_tok_s,
            # Output throughput per user (average per request output throughput)
            "system_total_throughput_tok_s":
            self.total_token_throughput_tok_s,
            "output_throughput_per_user_tok_s":
            self.per_user_output_throughput_tok_s,
            # Output throughput per GPU (total throughput / world size)
            "output_throughput_per_gpu_tok_s":
            self.output_throughput_tok_s / self.rt_cfg.mapping["world_size"],
            # Request latency percentiles
            "request_latency_percentiles_ms":
            self.statistics.request_latency_percentiles.model_dump(
                exclude_none=True, by_alias=True, mode='json') | {
                    k: self.convert_to_ms(v)
                    for k, v in self.statistics.request_latency_percentiles.
                    model_dump().items()
                },
        }

        if self.statistics.total_energy is not None:
            stats_dict["energy"] = {
                "total_energy_j":
                self.statistics.total_energy,
                "output_tps_per_w":
                self.statistics.output_tps_per_w,
                "average_gpu_power":
                self.statistics.total_gpu_power /
                self.rt_cfg.mapping["world_size"]
            }

        if self.streaming:
            avg_tpot = self.convert_to_ms(
                self.statistics.per_user_time_per_output_token_ns)

            stats_dict["streaming_metrics"] = {
                # NOTE: Excludes TTFT by nature as this is a genphase calculation.
                "token_output_speed_tok_s":
                self.per_user_generation_token_throughput_s,
                # Average per request time-to-first-token (TTFT)
                "avg_ttft_ms":
                self.convert_to_ms(
                    self.statistics.per_user_time_to_first_token_ns),
                # Average per request token time-per-output-token (TPOT)
                "avg_tpot_ms":
                avg_tpot,
                # Average per request Time-per-output-token percentiles (TPOT)
                "tpot_percentiles":
                self.statistics.tpot_percentiles.model_dump(
                    exclude_none=True, by_alias=True, mode='json') | {
                        k: self.convert_to_ms(v)
                        for k, v in
                        self.statistics.tpot_percentiles.model_dump().items()
                    },
                # Per request Time-to-first-token percentiles (TTFT)
                "ttft_percentiles":
                self.statistics.ttft_percentiles.model_dump(
                    exclude_none=True, by_alias=True, mode='json') | {
                        k: self.convert_to_ms(v)
                        for k, v in
                        self.statistics.ttft_percentiles.model_dump().items()
                    },
                "gen_tps_percentiles":
                self.statistics.generation_tp_percentiles.model_dump(
                    exclude_none=True, by_alias=True, mode='json') | {
                        k: self.convert_rate_to_s(v)
                        for k, v in self.statistics.generation_tp_percentiles.
                        model_dump().items()
                    },
            }

        spec_decoding, decoding_mode = False, None
        if (self.rt_cfg.decoding_config
                and self.rt_cfg.decoding_config.decoding_mode
                != SpeculativeDecodingMode.NONE):
            # cpp decoding
            spec_decoding = True
            decoding_mode = self.rt_cfg.decoding_config.decoding_mode.values[1]
        elif ("speculative_config" in self.kwargs
              and self.kwargs["speculative_config"] is not None):
            # pytorch speculative decoding
            spec_decoding = True
            spec_config = self.kwargs["speculative_config"]
            # Handle both dict (from YAML) and object types
            if isinstance(spec_config, dict):
                decoding_mode = spec_config.get("decoding_type")
            else:
                decoding_mode = spec_config.decoding_type
        if (spec_decoding):
            stats_dict["decoding_stats"] = {
                "mode":
                decoding_mode,
                "num_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_draft_tokens_percentiles":
                self.statistics.num_accepted_draft_tokens_percentiles.
                model_dump(exclude_none=True, by_alias=True, mode='json') if
                self.statistics.num_accepted_draft_tokens_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,
                                                                 mode='json')

        return stats_dict

    def report_statistics(self) -> None:
        """Report internal statistics about benchmark.

        Returns:
            BenchmarkStatistics: Benchmark statistics for the provided keeper.
        """
        stats_dict = self.get_statistics_dict()
        engine = stats_dict["engine"]
        machine = stats_dict.get("machine")
        world_info = stats_dict["world_info"]
        requests = stats_dict["request_info"]
        perf = stats_dict["performance"]
        streaming = stats_dict.get("streaming_metrics")
        decoding = stats_dict.get("decoding_stats", None)

        backend_info = ""
        if self.rt_cfg.backend not in ('pytorch', '_autodeploy'):
            config_path = self.rt_cfg.engine_dir / "config.json"
            with open(config_path, "r") as config:
                engine_config = json.load(config)
            build_cfg = engine_config["build_config"]
            pretrain_cfg = engine_config["pretrained_config"]

            backend_info = (
                "\n\n===========================================================\n"
                "= ENGINE DETAILS\n"
                "===========================================================\n"
                f"Model:\t\t\t{engine['model']}\n"
                f"Model Path:\t\t{engine['model_path']}\n"
                f"Revision:\t\t{engine['revision'] or 'N/A'}\n"
                f"Engine Directory:\t{engine['engine_dir']}\n"
                f"TensorRT LLM Version:\t{engine['version']}\n"
                f"Dtype:\t\t\t{pretrain_cfg['dtype']}\n"
                f"KV Cache Dtype:\t\t{pretrain_cfg['quantization']['kv_cache_quant_algo']}\n"
                f"Quantization:\t\t{pretrain_cfg['quantization']['quant_algo']}\n"
                f"Max Input Length:\t{build_cfg['max_input_len']}\n"
                f"Max Sequence Length:\t{build_cfg['max_seq_len']}\n"
                f"\n")
        else:
            backend_info = (
                "\n\n===========================================================\n"
                f"= {self.rt_cfg.backend.upper()} BACKEND\n"
                "===========================================================\n"
                f"Model:\t\t\t{engine['model']}\n"
                f"Model Path:\t\t{engine['model_path']}\n"
                f"Revision:\t\t{engine['revision'] or 'N/A'}\n"
                f"TensorRT LLM Version:\t{engine['version']}\n"
                f"Dtype:\t\t\t{engine['dtype']}\n"
                f"KV Cache Dtype:\t\t{engine['kv_cache_dtype']}\n"
                f"Quantization:\t\t{engine['quantization']}\n"
                # TODO
                # f"Max Input Length:\t{build_cfg['max_input_len']}\n"
                # f"Max Sequence Length:\t{build_cfg['max_seq_len']}\n"
                f"\n")

        kv_cache_percentage = world_info.get("kv_cache_percentage", None)
        if kv_cache_percentage is not None:
            kv_cache_percentage = f"{kv_cache_percentage * 100.0:.2f}%"

        machine_info = (
            "===========================================================\n"
            "= MACHINE DETAILS \n"
            "===========================================================\n")
        if machine is None:
            machine_info += "No GPU info available\n\n"
        else:
            name = machine.get("name", "Unknown")
            mem_total_str = f"{machine['memory.total']:.2f} GB" if machine.get(
                "memory.total") is not None else "N/A"
            mem_clock_str = f"{machine['clocks.mem']:.2f} GHz" if machine.get(
                'clocks.mem') is not None else "N/A"
            machine_info += f"{name}, memory {mem_total_str}, {mem_clock_str}\n\n"

        world_info = (
            "===========================================================\n"
            "= WORLD + RUNTIME INFORMATION \n"
            "===========================================================\n"
            f"TP Size:                {world_info['tp_size']}\n"
            f"PP Size:                {world_info['pp_size']}\n"
            f"EP Size:                {world_info['ep_size']}\n"
            f"Max Runtime Batch Size: {world_info['max_batch_size']}\n"
            f"Max Runtime Tokens:     {world_info['max_num_tokens']}\n"
            f"Scheduling Policy:      {world_info['scheduling_policy']}\n"
            f"KV Memory Percentage:   {kv_cache_percentage}\n"
            f"Issue Rate (req/sec):   {world_info['issue_rate']:.4E}\n"
            f"\n")

        req_lat_info = "\n".join(
            f"[Latency] {key.upper():<7}: {perf['request_latency_percentiles_ms'][key]:.4f}"
            for key in perf['request_latency_percentiles_ms'].keys())

        request_info = (
            "===========================================================\n"
            "= REQUEST DETAILS \n"
            "===========================================================\n"
            f"Number of requests:             {requests['num_requests']}\n"
            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"
        )

        perf_header = (
            "===========================================================\n"
            "= PERFORMANCE OVERVIEW \n"
            "===========================================================\n")

        perf_stats = (
            f"Request Throughput (req/sec):                     {perf['request_throughput_req_s']:.4f}\n"
            f"Total Output Throughput (tokens/sec):             {perf['system_output_throughput_tok_s']:.4f}\n"
            f"Total Token Throughput (tokens/sec):              {perf['system_total_throughput_tok_s']:.4f}\n"
            f"Total Latency (ms):                               {perf['total_latency_ms']:.4f}\n"
            f"Average request latency (ms):                     {perf['avg_request_latency_ms']:.4f}\n"
            # Output Throughput includes context/first token.
            f"Per User Output Throughput [w/ ctx] (tps/user):   {perf['output_throughput_per_user_tok_s']:.4f}\n"
            f"Per GPU Output Throughput (tps/gpu):              {perf['output_throughput_per_gpu_tok_s']:.4f}\n"
        )

        if streaming:
            streaming = stats_dict["streaming_metrics"]
            itl = streaming["tpot_percentiles"]
            ttft = streaming["ttft_percentiles"]

            tpot_stats = "\n".join(
                f"[TPOT] {key.upper():<7}: {itl[key]:.4f}" for key in
                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])

            ttft_stats = "\n".join(
                f"[TTFT] {key.upper():<7}: {ttft[key]:.4f}" for key in
                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])

            gen_tps_stats = "\n".join(
                f"[GTPS] {key.upper():<7}: {streaming['gen_tps_percentiles'][key]:.4f}"
                for key in
                ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"])

            perf_stats += (
                f"Average time-to-first-token [TTFT] (ms):          {streaming['avg_ttft_ms']:.4f}\n"
                f"Average time-per-output-token [TPOT] (ms):        {streaming['avg_tpot_ms']:.4f}\n"
                f"Per User Output Speed (tps/user):                 {streaming['token_output_speed_tok_s']:.4f}\n"
                "\n-- Per-Request Time-per-Output-Token [TPOT] Breakdown (ms)\n\n"
                f"{tpot_stats}\n"
                "\n-- Per-Request Time-to-First-Token [TTFT] Breakdown (ms) \n\n"
                f"{ttft_stats}\n"
                "\n-- Per-Request Generation Throughput [GTPS] Breakdown (tps/user)\n\n"
                f"{gen_tps_stats}\n")

        if "energy" in stats_dict:
            energy = stats_dict["energy"]
            perf_stats += (
                "\n-- Energy Metrics --------------------------------------\n\n"
                f"Total Energy (J):                                 {energy['total_energy_j']:.4f}\n"
                f"Output Tokens per Second per Watt (tps/W):         {energy['output_tps_per_w']:.4f}\n"
                f"Average GPU Power (W):                            {energy['average_gpu_power']:.4f}\n"
            )

        perf_stats += (
            "\n-- Request Latency Breakdown (ms) -----------------------\n\n"
            f"{req_lat_info}\n")

        decoding_stats = ""
        if decoding is not None:
            decoding = stats_dict["decoding_stats"]
            if self.get_max_draft_len() > 0:
                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_draft_tokens_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"
                    "-- Number of Draft Tokens Details --------------------------------\n\n"
                    "\n"
                    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"
                )

        logging_info = (f"{backend_info}"
                        f"{machine_info}"
                        f"{request_info}"
                        f"{world_info}"
                        f"{perf_header}"
                        f"{perf_stats}"
                        f"{decoding_stats}"
                        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 speculative_config."""
        # Try to get from speculative_config
        if ("speculative_config" in self.kwargs
                and self.kwargs["speculative_config"] is not None):
            spec_config = self.kwargs["speculative_config"]
            # Handle both dict (from YAML) and object types
            if isinstance(spec_config, dict):
                draft_len = (spec_config.get("max_draft_len")
                             or spec_config.get("num_nextn_predict_layers"))
                return draft_len or 0
            return spec_config.max_draft_len or 0

        return 0