metrics.py 7.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import time
5
6
from dataclasses import dataclass, field
from typing import Optional
7
8

import numpy as np
9
import prometheus_client
10

11
from vllm.config import SpeculativeConfig
12
13
14
15
16
17
18
from vllm.logger import init_logger

logger = init_logger(__name__)


@dataclass
class SpecDecodingStats:
19
20
21
22
23
24
25
26
27
    """Per-step iteration decoding stats from scheduler.

    Each scheduler step, statistics on spec decoding performance are
    aggregated across requests by the scheduler and returned to the
    frontend in EngineCoreOutputs->SchedulerStats.
    """

    num_spec_tokens: int
    num_drafts: int = 0
28
29
    num_draft_tokens: int = 0
    num_accepted_tokens: int = 0
30
    num_accepted_tokens_per_pos: list[int] = field(default_factory=list)
31

32
33
    @classmethod
    def new(cls, num_spec_tokens: int) -> "SpecDecodingStats":
34
35
36
37
        return cls(
            num_spec_tokens=num_spec_tokens,
            num_accepted_tokens_per_pos=[0] * num_spec_tokens,
        )
38

39
40
    def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
        self.num_drafts += 1
41
42
        self.num_draft_tokens += num_draft_tokens
        self.num_accepted_tokens += num_accepted_tokens
43
44
45
46
        assert num_accepted_tokens <= self.num_spec_tokens
        for i in range(num_accepted_tokens):
            self.num_accepted_tokens_per_pos[i] += 1

47

48
49
class SpecDecodingLogging:
    """Aggregate and log spec decoding metrics.
50

51
52
53
54
    LoggingStatLogger aggregates per-iteration metrics over a set
    time interval using observe() and then logs them using log()
    before resetting to zero.
    """
55
56
57
58
59

    def __init__(self):
        self.reset()

    def reset(self):
60
        self.num_drafts: list[int] = []
61
62
        self.num_draft_tokens: list[int] = []
        self.num_accepted_tokens: list[int] = []
63
        self.accepted_tokens_per_pos_lists: list[list[int]] = []
64
        self.last_log_time = time.monotonic()
65
66

    def observe(self, spec_decoding_stats: SpecDecodingStats):
67
        self.num_drafts.append(spec_decoding_stats.num_drafts)
68
        self.num_draft_tokens.append(spec_decoding_stats.num_draft_tokens)
69
        self.num_accepted_tokens.append(spec_decoding_stats.num_accepted_tokens)
70
        self.accepted_tokens_per_pos_lists.append(
71
72
            spec_decoding_stats.num_accepted_tokens_per_pos
        )
73

74
    def log(self, log_fn=logger.info):
75
76
        if not self.num_drafts:
            return
77
        num_drafts = np.sum(self.num_drafts)
78
79
        num_draft_tokens = np.sum(self.num_draft_tokens)
        num_accepted_tokens = np.sum(self.num_accepted_tokens)
80
81
82
83
84
85
86
        draft_throughput = 0
        accepted_throughput = 0

        elapsed_time = time.monotonic() - self.last_log_time
        if elapsed_time > 0:
            draft_throughput = num_draft_tokens / elapsed_time
            accepted_throughput = num_accepted_tokens / elapsed_time
87

88
89
90
91
92
        draft_acceptance_rate = (
            num_accepted_tokens / num_draft_tokens * 100
            if num_draft_tokens > 0
            else float("nan")
        )
93
94
95

        # Conventionally, mean acceptance length includes the bonus token
        mean_acceptance_length = 1 + (num_accepted_tokens / num_drafts)
96
97
98
99

        pos_matrix = np.array(self.accepted_tokens_per_pos_lists)
        acceptance_rates = np.sum(pos_matrix, axis=0) / num_drafts
        rates_str = ", ".join(f"{p:.3f}" for p in acceptance_rates)
100

101
        log_fn(
102
            "SpecDecoding metrics: "
103
            "Mean acceptance length: %.2f, "
104
105
            "Accepted throughput: %.2f tokens/s, "
            "Drafted throughput: %.2f tokens/s, "
106
            "Accepted: %d tokens, "
107
            "Drafted: %d tokens, "
108
109
            "Per-position acceptance rate: %s, "
            "Avg Draft acceptance rate: %.1f%%",
110
            mean_acceptance_length,
111
112
            accepted_throughput,
            draft_throughput,
113
114
            num_accepted_tokens,
            num_draft_tokens,
115
            rates_str,
116
            draft_acceptance_rate,
117
        )
118
        self.reset()
119
120
121
122
123
124
125
126
127
128


class SpecDecodingProm:
    """Record spec decoding metrics in Prometheus.

    The acceptance rate can be calculated using a PromQL query:

      rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
      rate(vllm:spec_decode_num_draft_tokens_total[$interval])

129
130
    The mean acceptance length (conventionally including bonus tokens)
    can be calculated using:
131

132
      1 + (
133
      rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
134
      rate(vllm:spec_decode_num_drafts[$interval]))
135
136
137
138
139
140
141

    A per-position acceptance rate vector can be computed using

      vllm:spec_decode_num_accepted_tokens_per_pos[$interval] /
      vllm:spec_decode_num_drafts[$interval]
    """

142
143
144
145
146
147
    _counter_cls = prometheus_client.Counter

    def __init__(
        self,
        speculative_config: Optional[SpeculativeConfig],
        labelnames: list[str],
148
        per_engine_labelvalues: dict[int, list[str]],
149
    ):
150
151
152
153
        self.spec_decoding_enabled = speculative_config is not None
        if not self.spec_decoding_enabled:
            return

154
155
156
        counter_drafts = self._counter_cls(
            name="vllm:spec_decode_num_drafts",
            documentation="Number of spec decoding drafts.",
157
158
            labelnames=labelnames,
        )
159
        self.counter_spec_decode_num_drafts = make_per_engine(
160
161
            counter_drafts, per_engine_labelvalues
        )
162
163
164
165

        counter_draft_tokens = self._counter_cls(
            name="vllm:spec_decode_num_draft_tokens",
            documentation="Number of draft tokens.",
166
167
            labelnames=labelnames,
        )
168
        self.counter_spec_decode_num_draft_tokens = make_per_engine(
169
170
            counter_draft_tokens, per_engine_labelvalues
        )
171
172
173
174

        counter_accepted_tokens = self._counter_cls(
            name="vllm:spec_decode_num_accepted_tokens",
            documentation="Number of accepted tokens.",
175
176
            labelnames=labelnames,
        )
177
        self.counter_spec_decode_num_accepted_tokens = make_per_engine(
178
179
            counter_accepted_tokens, per_engine_labelvalues
        )
180
181

        assert speculative_config is not None
182
183
184
185
186
        num_spec_tokens = (
            speculative_config.num_speculative_tokens
            if self.spec_decoding_enabled
            else 0
        )
187
        pos_labelnames = labelnames + ["position"]
188
        base_counter = self._counter_cls(
189
190
            name="vllm:spec_decode_num_accepted_tokens_per_pos",
            documentation="Accepted tokens per draft position.",
191
192
            labelnames=pos_labelnames,
        )
193
        self.counter_spec_decode_num_accepted_tokens_per_pos: dict[
194
195
196
197
198
199
200
            int, list[prometheus_client.Counter]
        ] = {
            idx: [base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens)]
            for idx, lv in per_engine_labelvalues.items()
        }

    def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0):
201
202
        if not self.spec_decoding_enabled:
            return
203
        self.counter_spec_decode_num_drafts[engine_idx].inc(
204
205
            spec_decoding_stats.num_drafts
        )
206
        self.counter_spec_decode_num_draft_tokens[engine_idx].inc(
207
208
            spec_decoding_stats.num_draft_tokens
        )
209
        self.counter_spec_decode_num_accepted_tokens[engine_idx].inc(
210
211
            spec_decoding_stats.num_accepted_tokens
        )
212
        for pos, counter in enumerate(
213
214
            self.counter_spec_decode_num_accepted_tokens_per_pos[engine_idx]
        ):
215
            counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos])
216
217


218
219
220
def make_per_engine(
    counter: prometheus_client.Counter, per_engine_labelvalues: dict[int, list[str]]
):
221
222
223
224
225
    """Create a counter for each label value."""
    return {
        idx: counter.labels(*labelvalues)
        for idx, labelvalues in per_engine_labelvalues.items()
    }