test_async_scheduling.py 13.3 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from itertools import repeat
4
5
6
from typing import Any

import pytest
7
import torch._dynamo.config as dynamo_config
8

9
from tests.utils import large_gpu_mark, single_gpu_only
10
from vllm import SamplingParams
11
from vllm.logprobs import Logprob
12
from vllm.platforms import current_platform
13
from vllm.sampling_params import StructuredOutputsParams
14
from vllm.v1.metrics.reader import Metric
15
16
17
18
19

from ...conftest import VllmRunner
from ...models.utils import check_outputs_equal

MODEL = "Qwen/Qwen3-0.6B"
20
MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
21
22


23
24
25
26
27
28
29
30
first_prompt = (
    "The following numbers of the sequence "
    + ", ".join(str(i) for i in range(10))
    + " are:"
)
example_prompts = [first_prompt, "In one word, the capital of France is "] + [
    f"Tell me about the number {i}: " for i in range(32)
]
31

32
33
default_params = dict(
    temperature=0.0,  # greedy
34
35
36
    max_tokens=30,
    # spec decoding currently doesn't support min_tokens
    # min_tokens=28,
37
)
38

39

40
@single_gpu_only
41
42
43
44
45
46
47
48
def test_without_spec_decoding(
    sample_json_schema,
    monkeypatch: pytest.MonkeyPatch,
):
    """Test consistency of combos of async scheduling, preemption,
    uni/multiproc executor, prefill chunking."""
    struct_outputs = StructuredOutputsParams(json=sample_json_schema)
    test_sampling_params: list[dict[str, Any]] = [
49
50
        dict(),
        # dict(min_tokens=20),
51
52
        dict(presence_penalty=-1.0),
        dict(bad_words=["the", " the"]),
53
54
        dict(logprobs=2),
        dict(logprobs=2, presence_penalty=-1.0),
55
        dict(structured_outputs=struct_outputs),
56
57
58
59
60
61
62
63
        dict(
            structured_outputs=struct_outputs,
            logprobs=2,
        ),
        dict(
            structured_outputs=struct_outputs,
            presence_penalty=-1.0,
        ),
64
        dict(
65
            structured_outputs=struct_outputs,
66
67
68
            logprobs=2,
            presence_penalty=-1.0,
        ),
69
70
    ]

71
72
73
74
75
76
77
78
79
80
    # test_preemption, executor, async_scheduling,
    # spec_config, test_prefill_chunking
    test_configs = [
        (False, "mp", False, None, False),
        (True, "mp", False, None, True),
        (False, "mp", True, None, False),
        (False, "uni", True, None, False),
        (True, "mp", True, None, False),
        (True, "uni", True, None, False),
        (False, "mp", True, None, True),
81
82
        (True, "mp", True, None, True),
        (True, "uni", True, None, True),
83
84
    ]

85
86
87
88
89
90
91
92
93
94
95
96
    if current_platform.is_rocm():
        # On ROCm, Only test with structured_outputs (deterministic)
        # and skip chunk_prefill (more variable).
        test_configs = [
            cfg
            for cfg in test_configs
            if not cfg[4]  # skip chunk_prefill=True
        ]
        test_sampling_params = [
            p for p in test_sampling_params if p.get("structured_outputs") is not None
        ]

97
    run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)
98

99

100
101
@single_gpu_only
@large_gpu_mark(min_gb=16)
102
def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
103
104
105
106
107
108
    """Test consistency and acceptance rates with some different combos of
    preemption, executor, async scheduling, prefill chunking,
    spec decoding model length.
    """

    spec_config = {
109
        "method": "eagle3",
110
        "num_speculative_tokens": 2,
111
        "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
112
    }
113
    # Set small draft model len to force doesn't-fit-in-drafter case.
114
115
    spec_config_short = spec_config | {"max_model_len": 50}

116
117
    struct_outputs = StructuredOutputsParams(json=sample_json_schema)

118
119
    test_sampling_params = [
        dict(),
120
121
        dict(presence_penalty=-1.0),
        dict(bad_words=["the", " the"]),
122
        dict(logprobs=2),
123
        dict(logprobs=2, presence_penalty=-1.0),
124
125
126
127
        dict(structured_outputs=struct_outputs),
        dict(
            structured_outputs=struct_outputs,
            logprobs=2,
128
            presence_penalty=-1.0,
129
        ),
130
131
    ]

132
133
134
135
136
137
138
139
140
141
142
143
    # test_preemption, executor, async_scheduling,
    # spec_config, test_prefill_chunking
    test_configs = [
        (False, "mp", False, None, False),
        (False, "mp", False, spec_config, False),
        (True, "mp", False, spec_config, True),
        (True, "uni", False, spec_config_short, True),
        (False, "mp", True, spec_config, False),
        (True, "mp", True, spec_config, False),
        (False, "mp", True, spec_config_short, True),
        (True, "uni", True, spec_config, False),
        (True, "uni", True, spec_config_short, False),
144
145
        (True, "mp", True, spec_config, True),
        (True, "uni", True, spec_config_short, True),
146
147
    ]

148
149
150
151
152
153
154
155
    # On ROCm, use TRITON_ATTN + float32 for better numerical consistency
    run_tests(
        monkeypatch,
        MTP_MODEL,
        test_configs,
        test_sampling_params,
        is_testing_with_spec_decoding=True,
    )
156
157
158
159
160
161
162
163


@dynamo_config.patch(cache_size_limit=16)
def run_tests(
    monkeypatch: pytest.MonkeyPatch,
    model: str,
    test_configs: list[tuple],
    test_sampling_params: list[dict[str, Any]],
164
    is_testing_with_spec_decoding: bool = False,
165
166
167
168
):
    """Test consistency of combos of async scheduling, preemption,
    uni/multiproc executor with spec decoding."""

169
    # Determine attention config based on platform
170
    attention_config = {"backend": "FLEX_ATTENTION"}
171
172

    with monkeypatch.context() as m:
173
        # lock matmul precision to full FP32 (IEEE)
174
        m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
175
        # m.setenv("VLLM_BATCH_INVARIANT", "1")
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
        outputs: list[tuple[str, list, list]] = []
        for n, (
            test_preemption,
            executor,
            async_scheduling,
            spec_config,
            test_prefill_chunking,
        ) in enumerate(test_configs, 1):
            test_str = f"{n}/{len(test_configs)}"
            test_results = run_test(
                model,
                test_str,
                test_sampling_params,
                test_preemption,
                executor,
                async_scheduling,
                spec_config,
                test_prefill_chunking=test_prefill_chunking,
194
                is_testing_with_spec_decoding=is_testing_with_spec_decoding,
195
                attention_config=attention_config,
196
197
198
199
200
201
202
            )
            outputs.append(test_results)

    baseline_config, baseline_tests, _ = outputs[0]
    _, _, baseline_acceptances = next(
        (o for o in outputs if o[2] is not None), (None, None, None)
    )
203

204
205
206
207
208
209
210
211
212
213
214
215
216
    print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}")

    failure = None
    for test_config, test_outputs, test_acceptance_rates in outputs[1:]:
        for (base_outs, base_logprobs), base_acceptance_rate, (
            test_outs,
            test_logprobs,
        ), test_acceptance_rate, params in zip(
            baseline_tests,
            baseline_acceptances or repeat(None),
            test_outputs,
            test_acceptance_rates or repeat(None),
            test_sampling_params,
217
        ):
218
219
220
221
222
223
224
            try:
                check_outputs_equal(
                    outputs_0_lst=base_outs,
                    outputs_1_lst=test_outs,
                    name_0=f"baseline=[{baseline_config}], params={params}",
                    name_1=f"config=[{test_config}], params={params}",
                )
225

226
                assert _all_logprobs_match(base_logprobs, test_logprobs)
227
228
229
230
231
232

                if (
                    base_acceptance_rate is not None
                    and test_acceptance_rate is not None
                ):
                    if "spec_mml=None" in test_config:
233
234
235
236
237
238
239
240
                        # Preemption causes more variance in acceptance rates
                        if (
                            current_platform.is_rocm()
                            and "preemption=True" in test_config
                        ):
                            tolerance = 0.10
                        else:
                            tolerance = 0.05
241
                        assert (
242
243
                            test_acceptance_rate > base_acceptance_rate
                            or test_acceptance_rate
244
                            == pytest.approx(base_acceptance_rate, rel=tolerance)
245
246
247
                        )
                    else:
                        # Currently the reported acceptance rate is expected to be
248
                        # lower when we sometimes skip drafting altogether.
249
                        assert test_acceptance_rate > 0.1
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
                print(
                    f"PASSED: config=[{test_config}], params={params}"
                    f" accept_rate={test_acceptance_rate}"
                )
            except AssertionError as e:
                print(
                    f"FAILED: config=[{test_config}], params={params}"
                    f" accept_rate={test_acceptance_rate}"
                )
                if failure is None:
                    failure = e

    if failure is not None:
        raise failure


def run_test(
    model: str,
    test_str: str,
    sampling_param_tests: list[dict[str, Any]],
    test_preemption: bool,
    executor: str,
    async_scheduling: bool,
    spec_config: dict[str, Any] | None,
    test_prefill_chunking: bool,
275
    is_testing_with_spec_decoding: bool = False,
276
    attention_config: dict[str, Any] | None = None,
277
278
279
):
    spec_decoding = spec_config is not None
    cache_arg: dict[str, Any] = (
280
        # Force preemptions
281
282
283
284
285
286
287
288
289
290
291
292
293
294
        dict(num_gpu_blocks_override=32)
        if test_preemption
        else dict(gpu_memory_utilization=0.9)
    )
    spec_mml = (spec_config or {}).get("max_model_len")
    test_config = (
        f"executor={executor}, preemption={test_preemption}, "
        f"async_sched={async_scheduling}, "
        f"chunk_prefill={test_prefill_chunking}, "
        f"spec_decoding={spec_decoding}, spec_mml={spec_mml}"
    )
    print("-" * 80)
    print(f"---- TESTING {test_str}: {test_config}")
    print("-" * 80)
295

296
297
298
299
    with VllmRunner(
        model,
        max_model_len=512,
        enable_chunked_prefill=test_prefill_chunking,
300
        # Force prefill chunking
301
302
303
304
        max_num_batched_tokens=48 if test_prefill_chunking else None,
        # enforce_eager=True,
        async_scheduling=async_scheduling,
        distributed_executor_backend=executor,
305
        dtype="float32",
306
307
        speculative_config=spec_config,
        disable_log_stats=False,
308
        attention_config=attention_config,
309
310
311
312
313
314
315
316
317
318
        **cache_arg,
    ) as vllm_model:
        results = []
        acceptance_rates: list[float] | None = [] if spec_decoding else None
        for override_params in sampling_param_tests:
            metrics_before = vllm_model.llm.get_metrics()
            print(f"----------- RUNNING PARAMS: {override_params}")
            results.append(
                vllm_model.generate(
                    example_prompts,
319
                    sampling_params=SamplingParams(**default_params, **override_params),
320
321
                    return_logprobs=True,
                )
322
            )
323
324
325
326
327
328
329
330
            metrics_after = vllm_model.llm.get_metrics()
            if acceptance_rates is not None:
                acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
                acceptance_rates.append(acceptance_rate)
                print(f"ACCEPTANCE RATE {acceptance_rate}")

            if test_preemption:
                preemptions = _get_count(
331
                    metrics_before, metrics_after, "vllm:num_preemptions"
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
                )
                assert preemptions > 0, "preemption test had no preemptions"

    if len(results) > 1:
        # First check that the different parameter configs
        # actually result in different output.
        for (other_test_outs, other_test_logprobs), params in zip(
            results[1:], sampling_param_tests[1:]
        ):
            with pytest.raises(AssertionError):
                check_outputs_equal(
                    outputs_0_lst=results[0][0],
                    outputs_1_lst=other_test_outs,
                    name_0=f"baseline params={params}",
                    name_1=f"other params={params}",
                )
                assert _all_logprobs_match(results[0][1], other_test_logprobs)
349

350
    return test_config, results, acceptance_rates
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365


def _all_logprobs_match(req_a, req_b) -> bool:
    return (
        req_a == req_b
        or len(req_a) == len(req_b)
        and all(
            len(seq_a) == len(seq_b)
            and all(_logprobs_match(a, b) for a, b in zip(seq_a, seq_b))
            for seq_a, seq_b in zip(req_a, req_b)
        )
    )


def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
366
    rel_tol, abs_tol = 1e-3, 1e-6
367
368
369
370
371
372
373
374
375
    return (
        len(lps_a) == len(lps_b)
        and lps_a.keys() == lps_b.keys()
        and all(
            a.decoded_token == b.decoded_token
            and a.rank == b.rank
            and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol)
            for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
        )
376
    )
377
378
379
380
381
382
383
384
385
386
387
388


def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
    draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
    accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
    return accept / draft if draft > 0 else 0.0


def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
    before_val = next(m.value for m in before if m.name == name)
    after_val = next(m.value for m in after if m.name == name)
    return after_val - before_val