| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Which rows the PP sampled-token broadcast must carry.""" |
| 4 | |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import numpy as np |
| 8 | |
| 9 | from vllm.v1.worker.gpu import pp_utils |
| 10 | |
| 11 | |
| 12 | def _batch( |
| 13 | num_computed, |
| 14 | prefill_len, |
| 15 | num_scheduled, |
| 16 | is_prefilling=None, |
| 17 | max_seq_len=None, |
| 18 | ): |
| 19 | return Mock( |
| 20 | num_reqs=len(num_computed), |
| 21 | num_computed_tokens_np=np.array(num_computed, dtype=np.int32), |
| 22 | prefill_len_np=np.array(prefill_len, dtype=np.int32), |
| 23 | num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32), |
| 24 | is_prefilling_np=( |
| 25 | np.array(is_prefilling, dtype=np.bool_) |
| 26 | if is_prefilling is not None |
| 27 | else None |
| 28 | ), |
| 29 | max_seq_len_np=( |
| 30 | np.array(max_seq_len, dtype=np.int32) if max_seq_len is not None else None |
| 31 | ), |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | def test_excludes_non_final_prefill_chunks(): |
| 36 | """Unchanged behaviour: a chunk that does not finish its prefill is skipped.""" |
| 37 | # Row 0 is a middle prefill chunk and produces no sample; row 1 finishes its |
| 38 | # prefill this step and therefore does. |
| 39 | batch = _batch( |
| 40 | num_computed=[512, 1000], |
| 41 | prefill_len=[4096, 1004], |
| 42 | num_scheduled=[448, 4], |
| 43 | ) |
| 44 | |
| 45 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 46 | |
| 47 | assert mask is not None |
| 48 | assert mask.tolist() == [False, True] |
| 49 | |
| 50 | |
| 51 | def test_none_when_no_row_samples(): |
| 52 | """Unchanged behaviour: an all-prefill batch needs no broadcast at all.""" |
| 53 | batch = _batch( |
| 54 | num_computed=[0, 512], |
| 55 | prefill_len=[4096, 4096], |
| 56 | num_scheduled=[448, 448], |
| 57 | ) |
| 58 | |
| 59 | assert pp_utils.compute_need_sampled_mask(batch) is None |
| 60 | |
| 61 | |
| 62 | def test_keeps_decoding_request_past_its_length_cap(): |
| 63 | """A decoding request must never be dropped from the broadcast. |
| 64 | |
| 65 | Speculative decoding advances `num_computed_tokens` several tokens per step, |
| 66 | so it can overrun `prompt_len + max_tokens` while the scheduler is still |
| 67 | running the request. Predicting "this one is finishing" and skipping its |
| 68 | broadcast freezes the earlier pipeline stages' `last_sampled_tokens` and |
| 69 | `draft_tokens` while the last rank keeps advancing its own, and the stages |
| 70 | then diverge permanently. |
| 71 | """ |
| 72 | batch = _batch( |
| 73 | # 14176 computed tokens is already past this request's own |
| 74 | # prompt_len + max_tokens; the scheduler is still running it. |
| 75 | num_computed=[14176], |
| 76 | prefill_len=[12175], |
| 77 | num_scheduled=[8], |
| 78 | ) |
| 79 | |
| 80 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 81 | |
| 82 | assert mask is not None |
| 83 | assert mask.tolist() == [True] |
| 84 | |
| 85 | |
| 86 | def test_decode_row_ahead_of_a_prefill_chunk(): |
| 87 | """Row order does not matter: only whether the row finishes its prefill.""" |
| 88 | batch = _batch( |
| 89 | num_computed=[10, 512], |
| 90 | prefill_len=[8, 4096], |
| 91 | num_scheduled=[1, 448], |
| 92 | ) |
| 93 | |
| 94 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 95 | |
| 96 | assert mask is not None |
| 97 | assert mask.tolist() == [True, False] |
| 98 | |
| 99 | |
| 100 | def test_excludes_final_prefill_chunk_reaching_its_length_cap(): |
| 101 | """A max_tokens=1 final chunk hands its token off and leaves the engine. |
| 102 | |
| 103 | Disaggregated prefill submits max_tokens=1, so the request finishes with |
| 104 | the final chunk's single sample and no follow-up step on this engine |
| 105 | reads the broadcast payload. |
| 106 | """ |
| 107 | batch = _batch( |
| 108 | num_computed=[1000], |
| 109 | prefill_len=[1004], |
| 110 | num_scheduled=[4], |
| 111 | is_prefilling=[True], |
| 112 | max_seq_len=[1005], |
| 113 | ) |
| 114 | |
| 115 | assert pp_utils.compute_need_sampled_mask(batch) is None |
| 116 | |
| 117 | |
| 118 | def test_keeps_final_prefill_chunk_with_generation_budget_left(): |
| 119 | """The same final chunk, but the request keeps decoding here.""" |
| 120 | batch = _batch( |
| 121 | num_computed=[1000], |
| 122 | prefill_len=[1004], |
| 123 | num_scheduled=[4], |
| 124 | is_prefilling=[True], |
| 125 | max_seq_len=[1004 + 256], |
| 126 | ) |
| 127 | |
| 128 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 129 | |
| 130 | assert mask is not None |
| 131 | assert mask.tolist() == [True] |
| 132 | |
| 133 | |
| 134 | def test_length_cap_overrun_decode_row_is_kept(): |
| 135 | """Decoding rows are never excluded, even past their length cap. |
| 136 | |
| 137 | Speculative decoding can transiently overrun prompt_len + max_tokens |
| 138 | while the scheduler still runs the request; dropping the row would |
| 139 | desynchronize the PP stages. |
| 140 | """ |
| 141 | batch = _batch( |
| 142 | num_computed=[14176], |
| 143 | prefill_len=[12175], |
| 144 | num_scheduled=[8], |
| 145 | is_prefilling=[False], |
| 146 | max_seq_len=[12176], |
| 147 | ) |
| 148 | |
| 149 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 150 | |
| 151 | assert mask is not None |
| 152 | assert mask.tolist() == [True] |
| 153 | |
| 154 | |
| 155 | def test_finishing_prefill_row_does_not_drop_a_sibling_row(): |
| 156 | """Only finishing rows are excluded; a decoding sibling keeps its row.""" |
| 157 | batch = _batch( |
| 158 | num_computed=[1000, 10], |
| 159 | prefill_len=[1004, 8], |
| 160 | num_scheduled=[4, 1], |
| 161 | is_prefilling=[True, False], |
| 162 | max_seq_len=[1005, 1024], |
| 163 | ) |
| 164 | |
| 165 | mask = pp_utils.compute_need_sampled_mask(batch) |
| 166 | |
| 167 | assert mask is not None |
| 168 | assert mask.tolist() == [False, True] |