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

4
import math
5
import time
6
from typing import Optional
7
8
9

import pytest

10
11
12
13
14
15
16
from tests.v1.engine.utils import (
    NUM_PROMPT_LOGPROBS_UNDER_TEST,
    NUM_SAMPLE_LOGPROBS_UNDER_TEST,
    STOP_STRINGS,
    DummyOutputProcessorTestVectors,
    MockEngineCore,
)
17
from vllm import PoolingParams
18
from vllm.logprobs import PromptLogprobs, SampleLogprobs
19
from vllm.outputs import CompletionOutput, RequestOutput
20
from vllm.sampling_params import RequestOutputKind, SamplingParams
21
22
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.v1.engine import EngineCoreRequest
23
from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector
24
from vllm.v1.metrics.stats import IterationStats
25

26
27
28
29
30
31
32
33
34
35
36
37
38
39

def _ref_convert_id_to_token(
    tokenizer: AnyTokenizer,
    token_id: int,
) -> str:
    """Reference impl of logprobs detokenization.

    Args:
      tokenizer: tokenizer used by the model under test
      token_id: convert this token id

    Returns:
      String representation of input token id
    """
40
    return tokenizer.decode([token_id]) or ""
41
42
43


@pytest.mark.parametrize(
44
45
46
47
48
49
50
    "request_output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY]
)
def test_incremental_detokenization(
    request_output_kind: RequestOutputKind, dummy_test_vectors
):
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)
    engine_core = MockEngineCore(tokens_list=dummy_test_vectors.generation_tokens)
51
52
53

    # Make N requests.
    requests = [
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
        EngineCoreRequest(
            request_id=f"request-{idx}",
            prompt_token_ids=prompt_tokens,
            mm_features=None,
            eos_token_id=None,
            arrival_time=0,
            lora_request=None,
            cache_salt=None,
            data_parallel_rank=None,
            sampling_params=SamplingParams(
                skip_special_tokens=False,
                spaces_between_special_tokens=False,
                output_kind=request_output_kind,
                stop=[],
                include_stop_str_in_output=False,
            ),
            pooling_params=None,
        )
72
        for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
73
74
75
    ]

    # Add requests to the detokenizer.
76
77
    for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):
        output_processor.add_request(request, prompt)
78
79
80
81
82
83
84
85
86
87

    gen_strings = {}
    gen_tokens = {}
    while True:
        # Mock output from the EngineCore.
        outputs = engine_core.get_outputs()
        if len(outputs) == 0:
            break

        # Step the Detokenizer.
88
        processed_outputs = output_processor.process_outputs(outputs)
89
90
        request_outputs = processed_outputs.request_outputs
        requests_to_abort = processed_outputs.reqs_to_abort
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
        assert len(requests_to_abort) == 0

        # Update tracking.
        for request_output in request_outputs:
            request_id = request_output.request_id
            new_text = request_output.outputs[0].text
            new_tokens = request_output.outputs[0].token_ids
            if request_id not in gen_strings:
                gen_strings[request_id] = new_text
                gen_tokens[request_id] = new_tokens
            else:
                gen_strings[request_id] += new_text
                gen_tokens[request_id].extend(new_tokens)

    # Confirmed tracked values matches what we expected.
    for idx, (ref_gen_str, ref_gen_toks) in enumerate(
107
108
        zip(dummy_test_vectors.generation_strings, dummy_test_vectors.generation_tokens)
    ):
109
110
111
112
113
114
        gen_str = gen_strings[f"request-{idx}"]
        gen_toks = gen_tokens[f"request-{idx}"]

        assert gen_str == ref_gen_str, f"{gen_str=}, {ref_gen_str=}"
        assert gen_toks == ref_gen_toks, f"{gen_toks=}, {ref_gen_toks=}"

115
116
    assert output_processor.get_num_unfinished_requests() == 0
    assert not output_processor.has_unfinished_requests()
117
118


119
def _validate_logprobs(
120
121
122
123
    gen_tokens: dict[str, list[int]],
    gen_logprobs: dict[str, Optional[SampleLogprobs]],
    gen_prompt_logprobs: dict[str, Optional[PromptLogprobs]],
    gen_cumulative_logprob: dict[str, float],
124
    dtv: DummyOutputProcessorTestVectors,
125
    request_id_list: list[str],
126
127
128
129
130
131
132
133
134
135
136
137
138
    num_sample_logprobs: Optional[int],
    num_prompt_logprobs: Optional[int],
) -> None:
    for req_idx, req_id in enumerate(request_id_list):
        new_tokens = gen_tokens[req_id]
        logprobs = gen_logprobs[req_id]
        prompt_logprobs = gen_prompt_logprobs[req_id]
        cumulative_logprob = gen_cumulative_logprob[req_id]
        prompt_token_ids = dtv.prompt_tokens[req_idx]
        ref_logprobs = dtv.generation_logprobs[req_idx]
        ref_prompt_logprobs = dtv.prompt_logprobs[req_idx]
        if num_sample_logprobs is not None:
            # Validate sample logprobs
139
140
141
142
143
            assert logprobs is not None, (
                f"Request {req_id} requires sample"
                " logprobs but sample logprobs are"
                " None."
            )
144
145
146
147
148
149
150
151
152
153
            # Require num sampled tokens to match num
            # sampled logprobs - especially important
            # to check since the detokenizer can cause
            # a request to finish early due to a stop
            # string being hit
            num_new_tokens = len(new_tokens)
            len_sample_logprobs = len(logprobs)
            assert num_new_tokens == len_sample_logprobs, (
                f"Request {req_id} has {num_new_tokens}"
                " completion tokens but has"
154
155
                f" {len_sample_logprobs} sample logprobs."
            )
156
            ref_cumulative_logprob = 0.0
157
158
159
            for idx, (sampled_token, pos_logprob_dict) in enumerate(
                zip(new_tokens, logprobs)
            ):
160
161
162
163
                # Break out the reference log probability value &
                # logprob token id tensors associated with this
                # position in the completion. Also break out the
                # sampled token ranks
164
165
166
                (ref_pos_logprob_toks, ref_pos_logprob_vals, ref_sampled_token_rank) = (
                    ref_logprobs[idx]
                )
167
168
169
170
171
                # For each position in the completion sequence,
                # ensure the actual sampled token is among the
                # logprobs
                assert sampled_token in pos_logprob_dict, (
                    f"Sampled token {sampled_token} not"
172
173
                    f" present in logprob at index {idx}"
                )
174
175
176

                # Validate number of sample logprobs
                num_lp_toks = len(pos_logprob_dict)
177
178
179
180
181
182
183
184
185
186
187
                assert (
                    num_lp_toks == num_sample_logprobs
                    or num_lp_toks == num_sample_logprobs + 1
                ), (
                    "Valid numbers of sample logprobs are"
                    f" {num_sample_logprobs} or"
                    f" {num_sample_logprobs + 1} but"
                    f" {num_lp_toks} logprobs found at"
                    f" position {idx}. Logprobs dict:"
                    f" {pos_logprob_dict}"
                )
188
189
190
191

                # Validate sampled token logprob rank
                smp_lp = pos_logprob_dict[sampled_token]
                smp_lp_rank = smp_lp.rank
192
                assert ref_sampled_token_rank == smp_lp_rank, (
193
194
195
196
                    "Sampled token logprob rank"
                    f" {smp_lp_rank} does not match"
                    " correct value"
                    f" {ref_sampled_token_rank}"
197
198
                    f" in Logprob {smp_lp}"
                )
199
200
201
202
203
204
205
206
207
208
209
210
211

                # Validate that the logprob processor yields
                # the correct log probabilities and valid
                # rankings
                rank_one_appears = False
                for jdx in range(1, len(ref_pos_logprob_toks)):
                    # Iterate over the (logprob val,logprob tok id)
                    # pairs expected by the test fixture at this
                    # position in the completion.
                    ref_lp_val = ref_pos_logprob_vals[jdx]
                    ref_tok_id = ref_pos_logprob_toks[jdx]
                    assert ref_tok_id in pos_logprob_dict, (
                        f"Expected token {ref_tok_id} to be"
212
213
                        f" in logprob dict but it is not."
                    )
214
215
216
217
218
219
220
221
222

                    # Extract actually-generated logprob
                    # info
                    lp = pos_logprob_dict[ref_tok_id]
                    lp_val = lp.logprob
                    lp_rank = lp.rank

                    # A "top" (rank 1) logprob must be
                    # present
223
                    rank_one_appears = True if lp_rank == 1 else rank_one_appears
224
225

                    # Rank must be >= 1
226
227
228
229
230
                    assert lp_rank >= 1, (
                        f"Logprob {lp} has invalid"
                        f" rank {lp_rank} < 1."
                        f" Logprob dict: {pos_logprob_dict}"
                    )
231
232
233
234
235
236

                    # Validate log probability
                    assert math.isclose(lp_val, ref_lp_val), (
                        f"Token id {ref_tok_id} appears in logprobs dict"
                        f" at position {idx} in completion with log"
                        f" probability {lp_val} but {ref_lp_val} was"
237
238
                        f" expected. Logprob: {lp}"
                    )
239

240
241
242
243
244
                assert rank_one_appears, (
                    f"No Logprob has rank 1"
                    " in the following Logprob"
                    f" dict: {pos_logprob_dict}"
                )
245
246
247
248
249
250

                # Validate logprobs detokenization
                for lp_tok in pos_logprob_dict:
                    # Confirm that sample logprob decoded token matches
                    # the logprob token id at this sequence position
                    decoded_token = pos_logprob_dict[lp_tok].decoded_token
251
                    ref_decoded_token = _ref_convert_id_to_token(dtv.tokenizer, lp_tok)
252
253
254
255
                    assert decoded_token == ref_decoded_token, (
                        f"Sampled logprob token id {lp_tok} decodes to"
                        f" {ref_decoded_token} but Logprob decoded"
                        f" token is {decoded_token} instead"
256
257
                        f" (at position {idx})"
                    )
258

259
                ref_cumulative_logprob += pos_logprob_dict[sampled_token].logprob
260
261
262
263
264
265
266
267
268
269
270
271
            # Assert that cumulative logprobs are correct
            assert math.isclose(cumulative_logprob, ref_cumulative_logprob)
        else:
            # Sample logprobs disabled for this request
            assert logprobs is None
            assert cumulative_logprob is None

        if num_prompt_logprobs is not None:
            # Validate prompt logprobs
            assert prompt_logprobs is not None, (
                f"Request {req_id} requires prompt"
                " logprobs but prompt logprobs are"
272
273
                " None."
            )
274
275
276
277
278
279
280
            # Require num prompt tokens to match num
            # prompt logprobs
            num_prompt_tokens = len(prompt_token_ids)
            len_prompt_logprobs = len(prompt_logprobs)
            assert num_prompt_tokens == len_prompt_logprobs, (
                f"Request {req_id} has {num_prompt_tokens}"
                " prompt tokens but has"
281
282
                f" {len_prompt_logprobs} prompt logprobs."
            )
283
284
285
286
287
            # First prompt logprob is None
            first_plp_dict = prompt_logprobs[0]
            assert first_plp_dict is None, (
                f"Request {req_id} first prompt logprob"
                f" should be None but has following value"
288
289
                f" instead: {first_plp_dict}"
            )
290
291
292
            # Break out the reference prompt log prob value &
            # logprob token id matrices for the whole prompt.
            # Also break out the prompt token rank vector
293
294
295
296
297
            (
                ref_prompt_logprob_toks,
                ref_prompt_logprob_vals,
                ref_prompt_token_ranks,
            ) = ref_prompt_logprobs
298
            for idx, (prompt_token, pos_logprob_dict) in enumerate(
299
300
                zip(prompt_token_ids[1:], prompt_logprobs[1:])
            ):
301
302
303
                # Break out the reference prompt log prob value
                # vector, prompt logprob token id vector, and
                # prompt token rank at the current position.
304
305
306
307
308
309
310
311
312
                (
                    ref_pos_prompt_logprob_toks,
                    ref_pos_prompt_logprob_vals,
                    ref_pos_prompt_token_rank,
                ) = (
                    ref_prompt_logprob_toks[idx, :],
                    ref_prompt_logprob_vals[idx, :],
                    ref_prompt_token_ranks[idx],
                )
313
314
315
316
317

                # For each position in the prompt sequence,
                # ensure the actual prompt token is among the
                # logprobs
                assert prompt_token in pos_logprob_dict, (
318
319
                    f"Prompt token {prompt_token} not present in logprob at index {idx}"
                )
320
321
                # Validate number of prompt logprobs
                num_plp_toks = len(pos_logprob_dict)
322
323
324
325
326
327
328
329
330
331
332
                assert (
                    num_plp_toks == num_prompt_logprobs
                    or num_plp_toks == num_prompt_logprobs + 1
                ), (
                    "Valid numbers of prompt logprobs are"
                    f" {num_prompt_logprobs} or"
                    f" {num_prompt_logprobs + 1} but"
                    f" {num_plp_toks} logprobs found at"
                    f" position {idx}. Logprobs dict:"
                    f" {pos_logprob_dict}"
                )
333
334
335
336
337

                # Validate prompt token logprob rank
                prmpt_tok_lp = pos_logprob_dict[prompt_token]
                prmpt_tok_lp_rank = prmpt_tok_lp.rank
                ref_prmpt_tok_lp_rank = ref_pos_prompt_token_rank
338
                assert ref_prmpt_tok_lp_rank == prmpt_tok_lp_rank, (
339
340
341
342
                    "Prompt token logprob rank"
                    f" {prmpt_tok_lp_rank} does not match"
                    " correct value"
                    f" {ref_prmpt_tok_lp_rank}"
343
344
                    f" in Logprob {prmpt_tok_lp}"
                )
345
346
347
348
349
350
351
352
353
354
355
356
357

                # Validate that the logprob processor yields
                # the correct prompt log probs and valid
                # rankings
                rank_one_appears = False
                for jdx in range(1, len(ref_pos_prompt_logprob_toks)):
                    # Iterate over the (logprob val,logprob tok id)
                    # pairs expected by the test fixture at this
                    # position in the completion.
                    ref_plp_val = float(ref_pos_prompt_logprob_vals[jdx])
                    ref_tok_id = int(ref_pos_prompt_logprob_toks[jdx])
                    assert ref_tok_id in pos_logprob_dict, (
                        f"Expected token {ref_tok_id} to be"
358
359
                        f" in logprob dict but it is not."
                    )
360
361
362
363
364
365
366
367
368

                    # Extract actually-generated logprob
                    # info
                    plp = pos_logprob_dict[ref_tok_id]
                    plp_val = plp.logprob
                    plp_rank = plp.rank

                    # A "top" (rank 1) logprob must be
                    # present
369
                    rank_one_appears = True if plp_rank == 1 else rank_one_appears
370
371
372
373
374

                    # Rank must be >= 1
                    assert plp_rank >= 1, (
                        f"Logprob {plp} has invalid"
                        f" rank {plp_rank} < 1."
375
376
                        f" Logprob dict: {pos_logprob_dict}"
                    )
377
378
379
380
381
382

                    # Validate log probability
                    assert math.isclose(plp_val, ref_plp_val), (
                        f"Token id {ref_tok_id} appears in logprobs dict"
                        f" at position {idx} in completion with log"
                        f" probability {plp_val} but {ref_plp_val} was"
383
384
                        f" expected. Logprob: {plp}"
                    )
385

386
387
388
389
390
                assert rank_one_appears, (
                    f"No Logprob has rank 1"
                    " in the following Logprob"
                    f" dict: {pos_logprob_dict}"
                )
391
392
393
394
395
396

                # Validate prompt logprob detokenization
                for plp_tok in pos_logprob_dict:
                    # Confirm that prompt logprob decoded token matches
                    # the logprob token id at this sequence position
                    decoded_token = pos_logprob_dict[plp_tok].decoded_token
397
                    ref_decoded_token = _ref_convert_id_to_token(dtv.tokenizer, plp_tok)
398
399
400
401
                    assert decoded_token == ref_decoded_token, (
                        f"Prompt logprob token id {plp_tok} decodes to"
                        f" {ref_decoded_token} but Logprob decoded"
                        f" token is {decoded_token} instead"
402
403
                        f" (at position {idx})"
                    )
404
405
406
407
408
409
        else:
            # Prompt logprobs disabled for this request
            assert prompt_logprobs is None


@pytest.mark.parametrize(
410
411
412
413
414
415
416
417
418
419
420
    "request_output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY]
)
@pytest.mark.parametrize("num_sample_logprobs", [None, NUM_SAMPLE_LOGPROBS_UNDER_TEST])
@pytest.mark.parametrize("num_prompt_logprobs", [None, NUM_PROMPT_LOGPROBS_UNDER_TEST])
def test_logprobs_processor(
    request_output_kind: RequestOutputKind,
    num_sample_logprobs: Optional[int],
    num_prompt_logprobs: Optional[int],
    dummy_test_vectors,
):
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)
421
422
    engine_core = MockEngineCore(
        tokens_list=dummy_test_vectors.generation_tokens,
423
424
425
        generated_logprobs_raw=None
        if num_sample_logprobs is None
        else dummy_test_vectors.generation_logprobs,
426
        prompt_logprobs_raw=None
427
428
429
        if num_prompt_logprobs is None
        else dummy_test_vectors.prompt_logprobs,
    )
430
431
432

    # Make N requests.
    request_id_list = [
433
        f"request-{idx}" for idx in range(len(dummy_test_vectors.prompt_strings))
434
435
    ]
    requests = [
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
        EngineCoreRequest(
            request_id=request_id_list[idx],
            prompt_token_ids=prompt_tokens,
            mm_features=None,
            eos_token_id=None,
            arrival_time=0,
            lora_request=None,
            cache_salt=None,
            data_parallel_rank=None,
            sampling_params=SamplingParams(
                skip_special_tokens=False,
                spaces_between_special_tokens=False,
                output_kind=request_output_kind,
                stop=[],
                include_stop_str_in_output=False,
                logprobs=num_sample_logprobs,
                prompt_logprobs=num_prompt_logprobs,
            ),
            pooling_params=None,
        )
456
        for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
457
458
459
    ]

    # Add requests to the detokenizer.
460
461
    for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):
        output_processor.add_request(request, prompt)
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485

    gen_tokens = {}
    gen_logprobs = {}
    gen_prompt_logprobs = {}
    gen_cumulative_logprobs = {}
    while True:
        # Mock output from the EngineCore.
        outputs = engine_core.get_outputs()
        if len(outputs) == 0:
            break

        # Step the logprobs processor.
        processed_outputs = output_processor.process_outputs(outputs)
        request_outputs = processed_outputs.request_outputs
        requests_to_abort = processed_outputs.reqs_to_abort
        assert len(requests_to_abort) == 0

        # Update tracking.
        for request_output in request_outputs:
            request_id = request_output.request_id
            new_tokens = request_output.outputs[0].token_ids
            prompt_logprobs = request_output.prompt_logprobs
            logprobs = request_output.outputs[0].logprobs
            gen_cumulative_logprobs[request_id] = request_output.outputs[
486
487
                0
            ].cumulative_logprob
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
            if request_id not in gen_logprobs:
                # Start tracking sample and prompt logprobs for this request
                gen_tokens[request_id] = new_tokens
                gen_logprobs[request_id] = logprobs
                gen_prompt_logprobs[request_id] = prompt_logprobs
            else:
                # Extend logprobs tracker
                gen_tokens[request_id].extend(new_tokens)
                lp = gen_logprobs[request_id]
                plp = gen_prompt_logprobs[request_id]
                if lp:
                    lp.extend(logprobs)
                if plp:
                    plp.extend(prompt_logprobs)

    # Confirmed tracked logprobs match what we expect
504
505
506
507
508
509
510
511
512
513
    _validate_logprobs(
        gen_tokens,
        gen_logprobs,
        gen_prompt_logprobs,
        gen_cumulative_logprobs,
        dummy_test_vectors,
        request_id_list,
        num_sample_logprobs,
        num_prompt_logprobs,
    )
514
515
516
517
518

    assert output_processor.get_num_unfinished_requests() == 0
    assert not output_processor.has_unfinished_requests()


519
520
@pytest.mark.parametrize(
    "include_stop_str_in_output,stop_token_type,ignore_eos,num_sample_logprobs",
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
    [
        (False, "stop_token_ids", False, None),
        (True, "stop_token_ids", False, None),
        (False, "stop_token_ids", False, NUM_SAMPLE_LOGPROBS_UNDER_TEST),
        (True, "stop_token_ids", False, NUM_SAMPLE_LOGPROBS_UNDER_TEST),
        (False, "eos_token_id", False, None),
        (True, "eos_token_id", False, None),
        (False, "eos_token_id", True, None),
    ],
)
def test_stop_token(
    include_stop_str_in_output: bool,
    num_sample_logprobs: Optional[int],
    stop_token_type: str,
    ignore_eos: bool,
    dummy_test_vectors,
):
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
    """Test output processor EOS/stop token handling.

    Send mock engine core request to mock engine core and pass core outputs
    to output processor. Validate output processor tokens, text and
    (if enabled) sample logprobs. Batch-size one.

    The test emulates a scenario where a model outputs text tokens followed
    by two identical control tokens:
    <token><token>...<token><control><control>

    If EOS is under test, the control tokens are EOS; otherwise, they are
    some other token id.

    Test behavior:

    * If EOS is under test and `ignore_eos=True`, the detokenized string
      should be <token><token>...<token><control><control> and the finish
      reason should be "length" (i.e. no stop occurs)

    * else, if `include_stop_str_in_output==True`, the detokenized
      string should be <token><token>...<token><control> and the finish
      reason should be "stop" (i.e. first control token causes stop
      and is represented in output text)

562
    * else, the detokenized string should be
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
      <token><token>...<token> and the finish reason should be "stop"
      (i.e. first control token causes stop but is not represented
      in output text.)

    Note: some test details are tuned for meta-llama/Llama-3.2-1B,
    another model should work only if the test is modified.

    Args:
        include_stop_str_in_output: stop token str appears in output text
        num_sample_logprobs: number of sample logprobs (`None` for no logprobs)
        stop_token_type: "eos_token_id" for EOS, "stop_token_ids" for stop token
        ignore_eos: if True, EOS stops are disabled
        dummy_test_vectors: dummy engine core outputs and other data structures
    """
    model_id = dummy_test_vectors.tokenizer.name_or_path
578
579
580
581
    if model_id != "meta-llama/Llama-3.2-1B":
        raise AssertionError(
            f"Test requires meta-llama/Llama-3.2-1B but {model_id} is in use."
        )
582
583
584
585
586
587
588
589
590
591
    do_logprobs = num_sample_logprobs is not None
    # EOS under test; if False, stop_token_ids under test
    is_eos_test = stop_token_type == "eos_token_id"
    # EOS under test but ignore_eos enabled
    is_eos_ignore_test = is_eos_test and ignore_eos
    eos_token_id = (
        dummy_test_vectors.tokenizer.eos_token_id if is_eos_test else None
    )  # '<|end_of_text|>'
    stop_token_ids = [128009] if not is_eos_test else None  # '<|eot_id|>'

592
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)
593
    # Dummy engine core outputs, with control tokens suffixed to test stops
594
    suffix_token = [eos_token_id] if is_eos_test else stop_token_ids
595
596
    assert suffix_token is not None and isinstance(suffix_token[0], int)
    generation_string = dummy_test_vectors.generation_strings[0]
597
    generation_tokens = dummy_test_vectors.generation_tokens[0] + 2 * suffix_token
598
    if do_logprobs:
599
600
601
        generation_logprobs = dummy_test_vectors.generation_logprobs[0] + 2 * [
            dummy_test_vectors.generation_logprobs[0][-1]
        ]
602
603
604
605
606
607
608
609
    prompt_string = dummy_test_vectors.prompt_strings[0]
    prompt_tokens = dummy_test_vectors.prompt_tokens[0]
    engine_core = MockEngineCore(
        tokens_list=[generation_tokens],
        generated_logprobs_raw=[generation_logprobs] if do_logprobs else None,
        prompt_logprobs_raw=None,
        eos_token_id=eos_token_id,
        stop_token_ids=stop_token_ids,
610
611
        ignore_eos=ignore_eos,
    )
612
613
614
615
616
617

    # Make request.
    request_id = "request-0"
    request = EngineCoreRequest(
        request_id=request_id,
        prompt_token_ids=prompt_tokens,
618
        mm_features=None,
619
        eos_token_id=eos_token_id,
620
        arrival_time=0,
621
        lora_request=None,
622
        cache_salt=None,
623
        data_parallel_rank=None,
624
625
626
627
628
629
630
631
632
633
        sampling_params=SamplingParams(
            skip_special_tokens=False,
            spaces_between_special_tokens=False,
            output_kind=RequestOutputKind.DELTA,
            stop=[],
            stop_token_ids=stop_token_ids,
            include_stop_str_in_output=include_stop_str_in_output,
            logprobs=num_sample_logprobs,
            prompt_logprobs=None,
            ignore_eos=ignore_eos,
634
        ),
635
636
        pooling_params=None,
    )
637
638

    # Add request to the detokenizer.
639
    output_processor.add_request(request, prompt_string)
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660

    # Loop over engine core steps; run output processor
    gen_string = ""
    gen_tokens = []
    gen_logprobs = []
    while True:
        # Mock output from the EngineCore.
        outputs = engine_core.get_outputs()
        if len(outputs) == 0:
            break

        # Step the Detokenizer.
        processed_outputs = output_processor.process_outputs(outputs)
        request_outputs = processed_outputs.request_outputs
        assert len(request_outputs) == 1
        # Stop token does not rely on abort
        assert not processed_outputs.reqs_to_abort

        # Update tracking.
        request_output = request_outputs[0]
        if request_output.finished:
661
            finish_reason = "length" if is_eos_ignore_test else "stop"
662
663
664
665
666
667
668
669
            assert request_output.outputs[0].finish_reason == finish_reason

        gen_string += request_output.outputs[0].text
        gen_tokens.extend(request_output.outputs[0].token_ids)
        if do_logprobs:
            gen_logprobs.extend(request_output.outputs[0].logprobs)

    # Validate generated text
670
    control_token = "<|end_of_text|>" if is_eos_test else "<|eot_id|>"
671
672
673
674
675
676
677
678
679
    if is_eos_ignore_test:
        # Length-based stop; expect full string
        ref_str = generation_string + 2 * control_token
    elif include_stop_str_in_output:
        # Stop token triggered; include in output
        ref_str = generation_string + control_token
    else:
        # Stop token triggered but not in output
        ref_str = generation_string
680
    assert gen_string == ref_str, f"{gen_string=}, {ref_str=}"
681
682
683
684
685
686

    if do_logprobs:
        # Validate number of sample logprobs
        num_tokens = len(gen_tokens)
        num_logprobs = len(gen_logprobs)
        assert num_tokens == num_logprobs, (
687
688
            f"Token count ({num_tokens}) != logprobs count ({num_logprobs})"
        )
689
690
691
692
693
694

    # Check requests are finished
    assert output_processor.get_num_unfinished_requests() == 0
    assert not output_processor.has_unfinished_requests()


695
@pytest.mark.parametrize("include_stop_str_in_output", [True, False])
696
697
698
699
700
701
702
@pytest.mark.parametrize("num_sample_logprobs", [None, NUM_SAMPLE_LOGPROBS_UNDER_TEST])
def test_stop_string(
    include_stop_str_in_output: bool,
    num_sample_logprobs: Optional[int],
    dummy_test_vectors,
):
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=False)
703
704
705
    engine_core = MockEngineCore(
        tokens_list=dummy_test_vectors.generation_tokens,
        generated_logprobs_raw=dummy_test_vectors.generation_logprobs
706
707
708
709
        if num_sample_logprobs
        else None,
        prompt_logprobs_raw=None,
    )
710
711

    # Make N requests.
712
    request_id_list = [
713
        f"request-{idx}" for idx in range(len(dummy_test_vectors.prompt_strings))
714
    ]
715
    requests = [
716
        EngineCoreRequest(
717
            request_id=request_id_list[idx],
718
            prompt_token_ids=prompt_tokens,
719
            mm_features=None,
720
            eos_token_id=None,
721
            arrival_time=0,
722
            lora_request=None,
723
            cache_salt=None,
724
            data_parallel_rank=None,
725
726
727
728
729
730
            sampling_params=SamplingParams(
                skip_special_tokens=False,
                spaces_between_special_tokens=False,
                output_kind=RequestOutputKind.DELTA,
                stop=STOP_STRINGS,
                include_stop_str_in_output=include_stop_str_in_output,
731
                logprobs=num_sample_logprobs,
732
                prompt_logprobs=None,
733
            ),
734
735
            pooling_params=None,
        )
736
        for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
737
738
739
    ]

    # Add requests to the detokenizer.
740
741
    for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):
        output_processor.add_request(request, prompt)
742
743

    gen_strings = {}
744
745
746
747
    gen_tokens = {}
    gen_logprobs = {}
    gen_prompt_logprobs = {}
    gen_cumulative_logprobs = {}
748
749
750
751
752
753
754
755
    aborted = []
    while True:
        # Mock output from the EngineCore.
        outputs = engine_core.get_outputs()
        if len(outputs) == 0:
            break

        # Step the Detokenizer.
756
757
758
        processed_outputs = output_processor.process_outputs(outputs)
        request_outputs = processed_outputs.request_outputs
        requests_to_abort = processed_outputs.reqs_to_abort
759
760
761
762
763
764
765
766
767
768
769
770
        for request_output in request_outputs:
            # If aborted, we should not get a request output.
            assert request_output.request_id not in aborted
        aborted.extend(requests_to_abort)

        # Update tracking.
        for request_output in request_outputs:
            if request_output.finished:
                assert request_output.outputs[0].finish_reason == "stop"

            request_id = request_output.request_id
            new_text = request_output.outputs[0].text
771
772
773
774
            new_tokens = request_output.outputs[0].token_ids
            prompt_logprobs = request_output.prompt_logprobs
            logprobs = request_output.outputs[0].logprobs
            gen_cumulative_logprobs[request_id] = request_output.outputs[
775
776
                0
            ].cumulative_logprob
777
778
            if request_id not in gen_strings:
                gen_strings[request_id] = new_text
779
780
781
                gen_tokens[request_id] = new_tokens
                gen_logprobs[request_id] = logprobs
                gen_prompt_logprobs[request_id] = prompt_logprobs
782
783
            else:
                gen_strings[request_id] += new_text
784
785
786
787
788
789
790
                gen_tokens[request_id].extend(new_tokens)
                lp = gen_logprobs[request_id]
                plp = gen_prompt_logprobs[request_id]
                if lp:
                    lp.extend(logprobs)
                if plp:
                    plp.extend(prompt_logprobs)
791
792

    # Confirmed tracked values matches what we expected.
793
    for idx, (ref_gen_str, stop_str) in enumerate(
794
795
        zip(dummy_test_vectors.generation_strings, STOP_STRINGS)
    ):
796
797
798
799
800
801
802
803
804
805
806
807
808
        # Request should be aborted.
        request_id = f"request-{idx}"
        assert request_id in aborted

        # Collected values that were generated.
        gen_str = gen_strings[request_id]

        # Construct reference strings.
        stop_str_idx = ref_gen_str.find(stop_str)
        ref_str_exc_stop = ref_gen_str[:stop_str_idx]
        ref_str_inc_stop = ref_gen_str[:stop_str_idx] + stop_str

        if include_stop_str_in_output:
809
            assert gen_str == ref_str_inc_stop, f"{gen_str=}, {ref_str_inc_stop=}"
810
        else:
811
            assert gen_str == ref_str_exc_stop, f"{gen_str=}, {ref_str_exc_stop=}"
812

813
    # Confirmed tracked logprobs match what we expect
814
815
816
817
818
819
820
821
822
823
    _validate_logprobs(
        gen_tokens,
        gen_logprobs,
        gen_prompt_logprobs,
        gen_cumulative_logprobs,
        dummy_test_vectors,
        request_id_list,
        num_sample_logprobs,
        None,
    )
824

825
826
827
828
    assert output_processor.get_num_unfinished_requests() == 0
    assert not output_processor.has_unfinished_requests()


829
def test_iteration_stats(dummy_test_vectors):
830
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=True)
831
    engine_core = MockEngineCore(dummy_test_vectors.generation_tokens)
832
    engine_core_timestamp = time.monotonic()
833
834
835
836
837
838

    # Make N requests.
    requests = [
        EngineCoreRequest(
            request_id=f"request-{idx}",
            prompt_token_ids=prompt_tokens,
839
            mm_features=None,
840
            eos_token_id=None,
841
            arrival_time=0,
842
            lora_request=None,
843
            cache_salt=None,
844
            data_parallel_rank=None,
845
            sampling_params=SamplingParams(),
846
            pooling_params=None,
847
848
        )
        for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
849
850
851
    ]

    # Add all requests except one to the OutputProcessor.
852
    num_active = len(dummy_test_vectors.generation_tokens) - 1
853
    for request in requests[:num_active]:
854
        output_processor.add_request(request, None)
855
856
857
858
    inactive_request = requests[num_active]

    # First iteration has 2 prefills.
    outputs = engine_core.get_outputs()[:num_active]
859
    iteration_stats = IterationStats()
860
861
862
863
864
865
866
    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
    total_prompt_tokens = sum(
        [
            len(prompt_tokens)
            for prompt_tokens in dummy_test_vectors.prompt_tokens[:num_active]
        ]
    )
867
868
869
870
871
872

    assert iteration_stats.num_prompt_tokens == total_prompt_tokens
    assert iteration_stats.num_generation_tokens == num_active

    # Just decodes in this step.
    outputs = engine_core.get_outputs()[:num_active]
873
    iteration_stats = IterationStats()
874
    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
875
876
877
878
879

    assert iteration_stats.num_prompt_tokens == 0
    assert iteration_stats.num_generation_tokens == num_active

    # Add a new request - prefill and 2 decodes in this step.
880
    output_processor.add_request(inactive_request, None)
881
882
    num_active += 1
    outputs = engine_core.get_outputs()[:num_active]
883
    iteration_stats = IterationStats()
884
    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
885
    total_prompt_tokens = len(dummy_test_vectors.prompt_tokens[num_active - 1])
886
887
888
889
890
891

    assert iteration_stats.num_prompt_tokens == total_prompt_tokens
    assert iteration_stats.num_generation_tokens == num_active

    # Just decodes in this step.
    outputs = engine_core.get_outputs()[:num_active]
892
    iteration_stats = IterationStats()
893
    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
894
895
896

    assert iteration_stats.num_prompt_tokens == 0
    assert iteration_stats.num_generation_tokens == num_active
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916


@pytest.mark.asyncio
async def test_request_output_collector():
    NUM_REQS = 3
    TEXT = "a"

    def make_outputs() -> list[RequestOutput]:
        return [
            RequestOutput(
                request_id="my-request-id",
                prompt=None,
                prompt_token_ids=[1, 2, 3],
                prompt_logprobs=None,
                outputs=[
                    CompletionOutput(
                        index=0,
                        text=TEXT,
                        token_ids=[idx],
                        cumulative_logprob=(idx + 1 * 1.0),
917
918
                        logprobs=[{"a": idx, "b": idx}],
                        finish_reason="length" if (idx == NUM_REQS - 1) else None,
919
920
921
                    )
                ],
                finished=(idx == NUM_REQS - 1),
922
923
            )
            for idx in range(NUM_REQS)
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
        ]

    collector = RequestOutputCollector(RequestOutputKind.DELTA)

    # CASE 1: Put then get.
    outputs = make_outputs()
    collector.put(outputs[0])
    output = await collector.get()
    assert not collector.ready.is_set()
    assert collector.output is None
    assert output.outputs[0].text == "a"
    assert output.outputs[0].token_ids == [0]

    # CASE 2: 2 puts then get.
    num_to_put = 2
    outputs = make_outputs()
    for i in range(num_to_put):
        collector.put(outputs[i])
    output = await collector.get()
    assert not collector.ready.is_set()
    assert collector.output is None

    assert not output.finished
    # Text, token_ids, and logprobs should get merged.
    assert output.outputs[0].text == TEXT * num_to_put
949
    for tok_0, tok_1 in zip(output.outputs[0].token_ids, list(range(num_to_put))):
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
        assert tok_0 == tok_1
    assert len(output.outputs[0].logprobs) == num_to_put

    # Cumulative logprobs should be the last one.
    cumulative_logprob_expected = 1.0 * num_to_put
    assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected

    # CASE 3: Put all 3 (including a finished).
    num_to_put = 3
    outputs = make_outputs()
    for i in range(num_to_put):
        collector.put(outputs[i])
    output = await collector.get()
    assert not collector.ready.is_set()
    assert collector.output is None

    assert output.finished
    assert output.outputs[0].finish_reason == "length"
    # Text, token_ids, and logprobs should get merged.
    assert output.outputs[0].text == TEXT * num_to_put
970
    for tok_0, tok_1 in zip(output.outputs[0].token_ids, list(range(num_to_put))):
971
972
973
974
975
976
        assert tok_0 == tok_1
    assert len(output.outputs[0].logprobs) == num_to_put

    # Cumulative logprobs should be the last one.
    cumulative_logprob_expected = 1.0 * num_to_put
    assert output.outputs[0].cumulative_logprob == cumulative_logprob_expected
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057


@pytest.mark.asyncio
async def test_cumulative_output_collector_n():
    """Test collector correctly handles multiple outputs by index."""
    collector = RequestOutputCollector(RequestOutputKind.CUMULATIVE)
    outputs = [
        RequestOutput(
            request_id="my-request-id",
            prompt=None,
            prompt_token_ids=[1, 2, 3],
            prompt_logprobs=None,
            outputs=[
                CompletionOutput(
                    index=0,
                    text="a",
                    token_ids=[0],
                    cumulative_logprob=None,
                    logprobs=None,
                    finish_reason=None,
                ),
                CompletionOutput(
                    index=1,
                    text="b",
                    token_ids=[1],
                    cumulative_logprob=None,
                    logprobs=None,
                    finish_reason=None,
                ),
            ],
            finished=False,
        ),
        RequestOutput(
            request_id="my-request-id",
            prompt=None,
            prompt_token_ids=[1, 2, 3],
            prompt_logprobs=None,
            outputs=[
                CompletionOutput(
                    index=0,
                    text="ab",
                    token_ids=[0, 1],
                    cumulative_logprob=None,
                    logprobs=None,
                    finish_reason=None,
                ),
                CompletionOutput(
                    index=2,
                    text="c",
                    token_ids=[2],
                    cumulative_logprob=None,
                    logprobs=None,
                    finish_reason=None,
                ),
            ],
            finished=False,
        ),
    ]
    for output in outputs:
        collector.put(output)

    # Get the output and check that the text and token_ids are correct.
    result = await collector.get()
    # We are expecting
    # [{index: 0, text: "ab"}, {index: 1, text: "b"}, {index: 2, text: "c"}]
    assert len(result.outputs) == 3
    # First is the one where index is 0
    first = [k for k in result.outputs if k.index == 0]
    assert len(first) == 1
    assert first[0].text == "ab"

    # Second is the one where index is 1
    second = [k for k in result.outputs if k.index == 1]
    assert len(second) == 1
    assert second[0].text == "b"
    assert second[0].token_ids == [1]

    # Third is the one where index is 2
    third = [k for k in result.outputs if k.index == 2]
    assert len(third) == 1
    assert third[0].text == "c"
1058
1059
1060
1061


@pytest.mark.parametrize("runner", ["generate", "pooling"])
def test_abort_requests(runner: str, dummy_test_vectors):
1062
    output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=True)
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
    requests = [
        EngineCoreRequest(
            request_id=f"request-{idx}",
            prompt_token_ids=prompt_tokens,
            mm_features=None,
            eos_token_id=None,
            arrival_time=0,
            lora_request=None,
            cache_salt=None,
            data_parallel_rank=None,
            sampling_params=SamplingParams() if runner == "generate" else None,
1074
1075
1076
            pooling_params=PoolingParams(task="embed") if runner == "pooling" else None,
        )
        for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
    ]

    for request in requests:
        if runner == "generate":
            output_kind = request.sampling_params.output_kind
        else:
            output_kind = request.pooling_params.output_kind
        queue = RequestOutputCollector(output_kind=output_kind)
        output_processor.add_request(request, None, queue=queue)

    for request in requests:
        output_processor.abort_requests([request.request_id])