| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for per-request speculative-decoding metrics accumulation. |
| 4 | |
| 5 | These cover the pure histogram/per-step math (no GPU / no model): the engine-core |
| 6 | accumulator ``RequestSpecDecodeMetrics`` and its ``to_dict`` payload surfaced in the |
| 7 | response as ``metrics.speculative_decoding``. |
| 8 | """ |
| 9 | |
| 10 | import msgspec |
| 11 | import pytest |
| 12 | |
| 13 | from vllm.outputs import CompletionOutput |
| 14 | from vllm.v1.engine import EngineCoreOutput |
| 15 | from vllm.v1.metrics.stats import RequestSpecDecodeMetrics |
| 16 | |
| 17 | |
| 18 | def _metrics(pairs, num_spec_tokens=3, detailed=False): |
| 19 | s = RequestSpecDecodeMetrics.new(num_spec_tokens) |
| 20 | for k, j in pairs: |
| 21 | s.observe(num_draft_tokens=k, num_accepted=j, detailed=detailed) |
| 22 | return s |
| 23 | |
| 24 | |
| 25 | def test_new_allocates_dense_histogram_of_k_plus_one(): |
| 26 | s = RequestSpecDecodeMetrics.new(num_spec_tokens=3) |
| 27 | assert s.num_spec_tokens == 3 |
| 28 | assert s.histogram == [0, 0, 0, 0] |
| 29 | assert s.num_draft_tokens == 0 |
| 30 | assert s.per_step_accepted == [] |
| 31 | |
| 32 | |
| 33 | def test_observe_buckets_by_accepted_draft_count(): |
| 34 | s = _metrics([(3, 0), (3, 3), (3, 2), (3, 3), (3, 1)]) |
| 35 | # j=0 ->1, j=1 ->1, j=2 ->1, j=3 ->2 |
| 36 | assert s.histogram == [1, 1, 1, 2] |
| 37 | assert s.num_draft_tokens == 15 |
| 38 | # summary level does not record the ordered per-step arrays |
| 39 | assert s.per_step_accepted == [] |
| 40 | assert s.per_step_drafted == [] |
| 41 | |
| 42 | |
| 43 | def test_observe_detailed_records_ordered_per_step_arrays(): |
| 44 | # Distinct step count (3), max draft length (k=4), and per-step drafted |
| 45 | # counts (4, 2, 4) so no accidental "everything is 3" pattern is implied. |
| 46 | s = _metrics([(4, 3), (2, 2), (4, 0)], num_spec_tokens=4, detailed=True) |
| 47 | assert s.per_step_accepted == [3, 2, 0] |
| 48 | assert s.per_step_drafted == [4, 2, 4] |
| 49 | # histogram (indexed by accepted j, length k+1=5) is still maintained |
| 50 | assert s.histogram == [1, 0, 1, 1, 0] |
| 51 | |
| 52 | |
| 53 | def test_to_dict_summary_omits_per_step_arrays(): |
| 54 | d = _metrics([(3, 0), (3, 3), (3, 2), (3, 3), (3, 1)]).to_dict() |
| 55 | assert d == { |
| 56 | "mean_acceptance_length": pytest.approx(1 + 9 / 5), # j+1 |
| 57 | "draft_acceptance_rate": pytest.approx(9 / 15), |
| 58 | "acceptance_histogram": [1, 1, 1, 2], # dense, index j = step count |
| 59 | "num_spec_steps": 5, |
| 60 | "num_accepted_draft_tokens": 9, |
| 61 | "num_draft_tokens": 15, |
| 62 | "num_spec_tokens": 3, |
| 63 | } |
| 64 | |
| 65 | |
| 66 | def test_to_dict_detailed_appends_per_step_arrays(): |
| 67 | d = _metrics([(3, 3), (3, 2), (3, 0)], detailed=True).to_dict() |
| 68 | assert d["per_step_accepted"] == [3, 2, 0] |
| 69 | assert d["per_step_drafted"] == [3, 3, 3] |
| 70 | # summary fields still present in detailed mode |
| 71 | assert d["num_spec_steps"] == 3 |
| 72 | assert d["num_accepted_draft_tokens"] == 5 |
| 73 | |
| 74 | |
| 75 | def test_to_dict_histogram_is_dense_list_indexed_by_j(): |
| 76 | # length k+1, index j holds the step count that accepted exactly j drafts |
| 77 | d = _metrics([(3, 0), (3, 0), (3, 3)]).to_dict() |
| 78 | assert d["acceptance_histogram"] == [2, 0, 0, 1] |
| 79 | |
| 80 | |
| 81 | def test_all_rejected_gives_mean_one_and_rate_zero(): |
| 82 | d = _metrics([(2, 0), (2, 0), (2, 0), (2, 0)], num_spec_tokens=2).to_dict() |
| 83 | assert d["num_spec_steps"] == 4 |
| 84 | assert d["num_accepted_draft_tokens"] == 0 |
| 85 | assert d["acceptance_histogram"] == [4, 0, 0] |
| 86 | assert d["mean_acceptance_length"] == pytest.approx(1.0) |
| 87 | assert d["draft_acceptance_rate"] == pytest.approx(0.0) |
| 88 | |
| 89 | |
| 90 | def test_empty_metrics_do_not_divide_by_zero(): |
| 91 | d = RequestSpecDecodeMetrics.new(3).to_dict() |
| 92 | assert d["num_spec_steps"] == 0 |
| 93 | assert d["num_draft_tokens"] == 0 |
| 94 | assert d["draft_acceptance_rate"] == 0.0 |
| 95 | assert d["mean_acceptance_length"] == 1.0 |
| 96 | assert "per_step_accepted" not in d |
| 97 | |
| 98 | |
| 99 | def test_observe_records_proposed_and_accepted_independently(): |
| 100 | # observe() takes proposed and accepted as independent inputs: the histogram |
| 101 | # is keyed by accepted, num_draft_tokens sums the proposed as given. (The |
| 102 | # grammar-invalidated-draft subtraction happens in the scheduler before |
| 103 | # observe() -- see test_per_request_spec_decode_subtracts_invalid_drafts.) |
| 104 | s = RequestSpecDecodeMetrics.new(num_spec_tokens=3) |
| 105 | s.observe(num_draft_tokens=2, num_accepted=1) |
| 106 | s.observe(num_draft_tokens=3, num_accepted=1) |
| 107 | d = s.to_dict() |
| 108 | assert d["acceptance_histogram"] == [0, 2, 0, 0] # both steps accepted 1 |
| 109 | assert d["num_draft_tokens"] == 5 # proposed summed independently: 2 + 3 |
| 110 | |
| 111 | |
| 112 | def test_engine_core_output_round_trips_spec_decode_metrics(): |
| 113 | # The accumulator rides EngineCoreOutput (msgspec, array_like) to the |
| 114 | # frontend; verify it serializes (incl. per-step arrays) and is omitted |
| 115 | # when absent. ``new_token_ids`` is EngineCoreOutput's required "tokens |
| 116 | # generated this step" field -- a dummy value here since we only exercise |
| 117 | # spec_decode_metrics. |
| 118 | metrics = _metrics([(3, 0), (3, 3), (3, 2)], detailed=True) |
| 119 | out = EngineCoreOutput( |
| 120 | request_id="r1", new_token_ids=[1, 2], spec_decode_metrics=metrics |
| 121 | ) |
| 122 | decoder = msgspec.msgpack.Decoder(EngineCoreOutput) |
| 123 | decoded = decoder.decode(msgspec.msgpack.encode(out)) |
| 124 | assert decoded.spec_decode_metrics.histogram == [1, 0, 1, 1] |
| 125 | assert decoded.spec_decode_metrics.num_draft_tokens == 9 |
| 126 | assert decoded.spec_decode_metrics.per_step_accepted == [0, 3, 2] |
| 127 | |
| 128 | without = EngineCoreOutput(request_id="r2", new_token_ids=[1]) |
| 129 | decoded_without = decoder.decode(msgspec.msgpack.encode(without)) |
| 130 | assert decoded_without.spec_decode_metrics is None |
| 131 | |
| 132 | |
| 133 | def _completion_output(**kwargs): |
| 134 | return CompletionOutput( |
| 135 | index=0, |
| 136 | text="", |
| 137 | token_ids=[], |
| 138 | cumulative_logprob=None, |
| 139 | logprobs=None, |
| 140 | **kwargs, |
| 141 | ) |
| 142 | |
| 143 | |
| 144 | def test_completion_output_carries_spec_decode_metrics(): |
| 145 | # Metrics are per output sequence, so they ride the CompletionOutput |
| 146 | # (=> choices[i]), not the request-level RequestOutput. |
| 147 | metrics = _metrics([(3, 3), (3, 1)]) |
| 148 | out = _completion_output(spec_decode_metrics=metrics) |
| 149 | assert out.spec_decode_metrics is metrics |
| 150 | assert _completion_output().spec_decode_metrics is None |