| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from unittest import TestCase |
| 4 | |
| 5 | import numpy as np |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.v1.outputs import ( |
| 11 | EMPTY_MODEL_RUNNER_OUTPUT, |
| 12 | ECConnectorOutput, |
| 13 | LogprobsLists, |
| 14 | LogprobsTensors, |
| 15 | ModelRunnerOutput, |
| 16 | SamplingMaskLists, |
| 17 | ) |
| 18 | from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p |
| 19 | from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors |
| 20 | |
| 21 | DEVICE_TYPE = current_platform.device_type |
| 22 | |
| 23 | |
| 24 | def test_logprobs_tensors_cat(): |
| 25 | first = LogprobsTensors( |
| 26 | torch.tensor([[1, 2]]), |
| 27 | torch.tensor([[0.1, 0.2]]), |
| 28 | torch.tensor([1]), |
| 29 | ) |
| 30 | second = LogprobsTensors( |
| 31 | torch.tensor([[3, 4]]), |
| 32 | torch.tensor([[0.3, 0.4]]), |
| 33 | torch.tensor([2]), |
| 34 | ) |
| 35 | |
| 36 | result = LogprobsTensors.cat([first, second], [0, 1, 2]) |
| 37 | |
| 38 | assert result.logprob_token_ids.tolist() == [[1, 2], [3, 4]] |
| 39 | assert result.logprobs.tolist() == ( |
| 40 | first.logprobs.tolist() + second.logprobs.tolist() |
| 41 | ) |
| 42 | assert result.selected_token_ranks.tolist() == [1, 2] |
| 43 | assert result.cu_num_generated_tokens == [0, 1, 2] |
| 44 | assert LogprobsTensors.cat([first]) is first |
| 45 | |
| 46 | |
| 47 | def test_logprobs_tensors_tolists_with_tensor_boundaries(): |
| 48 | """Adaptive verification hands over the request boundaries as a tensor |
| 49 | (they only exist on device); tolists() must materialize it as a plain |
| 50 | list so slice_request splits requests correctly.""" |
| 51 | tensors = LogprobsTensors( |
| 52 | torch.tensor([[1, 2], [3, 4], [5, 6]]), |
| 53 | torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), |
| 54 | torch.tensor([1, 2, 3]), |
| 55 | cu_num_generated_tokens_tensor=torch.tensor([0, 2, 3], dtype=torch.int32), |
| 56 | ) |
| 57 | |
| 58 | lists = tensors.to_cpu_nonblocking().tolists() |
| 59 | |
| 60 | assert lists.cu_num_generated_tokens == [0, 2, 3] |
| 61 | sliced = lists.slice_request(1, 1) |
| 62 | assert sliced.logprob_token_ids.tolist() == [[5, 6]] |
| 63 | assert sliced.sampled_token_ranks.tolist() == [3] |
| 64 | |
| 65 | |
| 66 | def test_sampling_mask_lists_to_nested_list(): |
| 67 | mask = SamplingMaskLists( |
| 68 | token_ids=np.array([10, 11, 12, 20, 21]), |
| 69 | offsets=np.array([0, 3, 5]), |
| 70 | ) |
| 71 | |
| 72 | nested = mask.to_nested_list() |
| 73 | |
| 74 | assert nested == [[10, 11, 12], [20, 21]] |
| 75 | assert SamplingMaskLists(np.array([7, 9])).to_nested_list() == [[7, 9]] |
| 76 | |
| 77 | |
| 78 | def test_sampling_mask_lists_slices_multiple_positions_by_request(): |
| 79 | masks = SamplingMaskLists( |
| 80 | token_ids=np.array([10, 11, 20, 30, 31, 32, 40]), |
| 81 | offsets=np.array([0, 2, 3, 6, 7]), |
| 82 | cu_num_generated_tokens=[0, 1, 3, 4], |
| 83 | ) |
| 84 | |
| 85 | single = masks.slice_request(0, 1) |
| 86 | multi = masks.slice_request(1, 2) |
| 87 | |
| 88 | assert single.to_nested_list() == [[10, 11]] |
| 89 | assert single.offsets is None |
| 90 | assert multi.to_nested_list() == [[20], [30, 31, 32]] |
| 91 | assert multi.cu_num_generated_tokens is None |
| 92 | |
| 93 | |
| 94 | @pytest.mark.parametrize("max_num_kept", [512, 20_001]) |
| 95 | @pytest.mark.skipif( |
| 96 | current_platform.is_xpu(), |
| 97 | reason="CI failed on XPU with Triton, Skipped until sycl kernel merged. " |
| 98 | "See https://github.com/vllm-project/vllm/pull/55851", |
| 99 | ) |
| 100 | def test_sampling_mask_tensors_match_finite_support(max_num_kept): |
| 101 | """Whatever its size (empty, one, around the compact width, beyond the |
| 102 | cap, the whole vocab), a sampled row's mask is exactly the ascending set |
| 103 | of its finite logits; rows that sampled nothing are empty.""" |
| 104 | from vllm.v1.worker.gpu.sample.output import MAX_COMPACT_SUPPORT |
| 105 | |
| 106 | vocab_size = 20_001 |
| 107 | sizes = [0, 1, 7, 511, 512, 513, MAX_COMPACT_SUPPORT + 5, vocab_size, 40, 3] |
| 108 | gen = torch.Generator().manual_seed(0) |
| 109 | logits = torch.full((len(sizes), vocab_size), float("-inf")) |
| 110 | expected = [] |
| 111 | for row, size in enumerate(sizes): |
| 112 | kept = torch.randperm(vocab_size, generator=gen)[:size].sort().values |
| 113 | logits[row, kept] = torch.randn(size, generator=gen) |
| 114 | expected.append(kept.tolist()) |
| 115 | num_sampled_tokens = torch.ones(len(sizes), dtype=torch.int32) |
| 116 | num_sampled_tokens[[1, 7]] = 0 |
| 117 | expected[1] = expected[7] = [] |
| 118 | |
| 119 | tensors = SamplingMaskTensors.from_logits( |
| 120 | logits.to(DEVICE_TYPE), |
| 121 | torch.arange(len(sizes) + 1, device=DEVICE_TYPE, dtype=torch.int32), |
| 122 | num_sampled_tokens.to(DEVICE_TYPE), |
| 123 | max_num_kept, |
| 124 | ) |
| 125 | |
| 126 | assert tensors.token_ids.shape[1] == min(max_num_kept, MAX_COMPACT_SUPPORT) |
| 127 | assert tensors.to_cpu_nonblocking().tolists().to_nested_list() == expected |
| 128 | |
| 129 | |
| 130 | def test_sampling_mask_matches_processed_top_k_top_p_support(): |
| 131 | """The mask must exactly mirror whatever support `apply_top_k_top_p` |
| 132 | (the real logits-processing function used by the sampler) actually |
| 133 | produces, whatever backend implements it.""" |
| 134 | processed_logits = apply_top_k_top_p( |
| 135 | logits=torch.tensor( |
| 136 | [[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE |
| 137 | ), |
| 138 | k=torch.tensor([3], device=DEVICE_TYPE), |
| 139 | p=None, |
| 140 | ) |
| 141 | expected_token_ids = ( |
| 142 | torch.isfinite(processed_logits[0]).nonzero().flatten().tolist() |
| 143 | ) |
| 144 | assert 0 < len(expected_token_ids) <= processed_logits.shape[1] |
| 145 | |
| 146 | tensors = SamplingMaskTensors.from_logits( |
| 147 | processed_logits, |
| 148 | cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE), |
| 149 | num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE), |
| 150 | max_num_kept=3, |
| 151 | ) |
| 152 | result = tensors.tolists() |
| 153 | |
| 154 | assert result.to_nested_list() == [expected_token_ids] |
| 155 | |
| 156 | |
| 157 | def test_sampling_mask_tensors_uses_request_boundaries(): |
| 158 | logits = torch.full((6, 8), -float("inf"), device=DEVICE_TYPE) |
| 159 | logits[0, [1, 3]] = 0 |
| 160 | logits[2, [2, 4, 6]] = 0 |
| 161 | logits[5, [0, 7]] = 0 |
| 162 | counts = torch.tensor([1, 0, 1], device=DEVICE_TYPE, dtype=torch.int32) |
| 163 | |
| 164 | tensors = SamplingMaskTensors.from_logits( |
| 165 | logits, |
| 166 | torch.tensor([0, 2, 5, 6], device=DEVICE_TYPE, dtype=torch.int32), |
| 167 | counts, |
| 168 | max_num_kept=3, |
| 169 | ) |
| 170 | |
| 171 | assert tensors.tolists().to_nested_list() == [[1, 3], [], [0, 7]] |
| 172 | |
| 173 | |
| 174 | def test_sampling_mask_tensors_multirow_request_layout(): |
| 175 | rows_per_request = 4 |
| 176 | counts = np.array([1, 2, 4, 0], dtype=np.int32) |
| 177 | logits = torch.full((16, 16), -float("inf"), device=DEVICE_TYPE) |
| 178 | expected = [] |
| 179 | for row in [0, 4, 5, 8, 9, 10, 11]: |
| 180 | kept = torch.tensor( |
| 181 | [row % 16, (row + 3) % 16, (row + 7) % 16], device=DEVICE_TYPE |
| 182 | ).unique(sorted=True) |
| 183 | logits[row, kept] = 0 |
| 184 | expected.append(kept.tolist()) |
| 185 | |
| 186 | tensors = SamplingMaskTensors.from_logits( |
| 187 | logits, |
| 188 | torch.arange(0, 17, 4, device=DEVICE_TYPE, dtype=torch.int32), |
| 189 | torch.from_numpy(counts).to(DEVICE_TYPE), |
| 190 | max_num_kept=2, |
| 191 | rows_per_request=rows_per_request, |
| 192 | ) |
| 193 | result = tensors.to_cpu_nonblocking().tolists(counts) |
| 194 | |
| 195 | assert tensors.token_ids.shape == (16, 2) |
| 196 | assert tensors.rows_per_request == rows_per_request |
| 197 | assert result.to_nested_list() == expected |
| 198 | assert result.cu_num_generated_tokens == [0, 1, 3, 7, 7] |
| 199 | |
| 200 | |
| 201 | def test_sampling_mask_preserves_top_k_boundary_ties(): |
| 202 | """When the kept support is wider than `max_num_kept` (e.g. a top-k |
| 203 | boundary tie keeps more than k logits), the mask must fall back to the |
| 204 | exact bitmask instead of silently truncating to `max_num_kept` ids.""" |
| 205 | processed_logits = torch.tensor( |
| 206 | [[6.0, 5.0, 4.0, 4.0, 4.0, float("-inf"), float("-inf"), float("-inf")]], |
| 207 | device=DEVICE_TYPE, |
| 208 | ) |
| 209 | expected_token_ids = [0, 1, 2, 3, 4] |
| 210 | |
| 211 | tensors = SamplingMaskTensors.from_logits( |
| 212 | processed_logits, |
| 213 | cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE), |
| 214 | num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE), |
| 215 | max_num_kept=3, |
| 216 | ) |
| 217 | result = tensors.tolists() |
| 218 | |
| 219 | assert result.to_nested_list() == [expected_token_ids] |
| 220 | |
| 221 | |
| 222 | class TestLogprobsLists(TestCase): |
| 223 | def setUp(self): |
| 224 | self.logprobsLists = LogprobsLists( |
| 225 | logprob_token_ids=[ |
| 226 | [1, 2], # Request 0 token 0 |
| 227 | [3, 4], # Request 0 token 1 |
| 228 | [5, 6], # Request 1 token 0 |
| 229 | [7, 8], # Request 1 token 1 |
| 230 | [9, 10], # Request 1 token 2 |
| 231 | [11, 12], # Request 2 token 0 |
| 232 | [13, 14], # Request 2 token 1 |
| 233 | [15, 16], # Request 2 token 2 |
| 234 | [17, 18], # Request 2 token 3 |
| 235 | ], |
| 236 | logprobs=[ |
| 237 | [0.1, 0.2], |
| 238 | [0.3, 0.4], |
| 239 | [0.5, 0.6], |
| 240 | [0.7, 0.8], |
| 241 | [0.9, 1.0], |
| 242 | [1.1, 1.2], |
| 243 | [1.3, 1.4], |
| 244 | [1.5, 1.6], |
| 245 | [1.7, 1.8], |
| 246 | ], |
| 247 | sampled_token_ranks=[1, 3, 5, 7, 9, 11, 13, 15, 17], |
| 248 | cu_num_generated_tokens=[0, 2, 5, 9], |
| 249 | ) |
| 250 | |
| 251 | def test_slice_without_cu_num_generated_tokens(self): |
| 252 | """Test slicing without cu_num_generated_tokens.""" |
| 253 | logprobsLists = LogprobsLists( |
| 254 | logprob_token_ids=[[1], [2], [3]], |
| 255 | logprobs=[[0.1], [0.2], [0.3]], |
| 256 | sampled_token_ranks=[1, 2, 3], |
| 257 | cu_num_generated_tokens=None, |
| 258 | ) |
| 259 | |
| 260 | sliced = logprobsLists.slice_request(1, num_positions=2) |
| 261 | assert sliced.logprob_token_ids == [[2], [3]] |
| 262 | assert sliced.logprobs == [[0.2], [0.3]] |
| 263 | assert sliced.sampled_token_ranks == [2, 3] |
| 264 | assert sliced.cu_num_generated_tokens is None |
| 265 | |
| 266 | def test_slice_from_start(self): |
| 267 | """Test slicing from the start position.""" |
| 268 | sliced = self.logprobsLists.slice_request(0, num_positions=5) |
| 269 | assert len(sliced.logprob_token_ids) == 5 |
| 270 | assert sliced.logprob_token_ids == [ |
| 271 | [1, 2], |
| 272 | [3, 4], |
| 273 | [5, 6], |
| 274 | [7, 8], |
| 275 | [9, 10], |
| 276 | ] |
| 277 | assert sliced.cu_num_generated_tokens is None |
| 278 | |
| 279 | def test_slice_from_middle(self): |
| 280 | """Test slicing from the middle position.""" |
| 281 | sliced = self.logprobsLists.slice_request(1, num_positions=7) |
| 282 | assert len(sliced.logprob_token_ids) == 7 |
| 283 | assert sliced.logprob_token_ids == [ |
| 284 | [5, 6], |
| 285 | [7, 8], |
| 286 | [9, 10], |
| 287 | [11, 12], |
| 288 | [13, 14], |
| 289 | [15, 16], |
| 290 | [17, 18], |
| 291 | ] |
| 292 | assert sliced.cu_num_generated_tokens is None |
| 293 | |
| 294 | def test_slice_single_request(self): |
| 295 | """Test slicing a single request.""" |
| 296 | sliced = self.logprobsLists.slice_request(1, num_positions=3) |
| 297 | assert len(sliced.logprob_token_ids) == 3 |
| 298 | assert sliced.logprob_token_ids == [[5, 6], [7, 8], [9, 10]] |
| 299 | assert sliced.cu_num_generated_tokens is None |
| 300 | |
| 301 | def test_slice_last_request(self): |
| 302 | """Test slicing the last request.""" |
| 303 | sliced = self.logprobsLists.slice_request(2, num_positions=4) |
| 304 | assert len(sliced.logprob_token_ids) == 4 |
| 305 | assert sliced.logprob_token_ids == [[11, 12], [13, 14], [15, 16], [17, 18]] |
| 306 | assert sliced.cu_num_generated_tokens is None |
| 307 | |
| 308 | def test_slice_all_requests(self): |
| 309 | """Test slicing all requests (full slice).""" |
| 310 | sliced = self.logprobsLists.slice_request(0, num_positions=9) |
| 311 | assert len(sliced.logprob_token_ids) == 9 # All tokens |
| 312 | assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids |
| 313 | assert sliced.cu_num_generated_tokens is None |
| 314 | |
| 315 | |
| 316 | def test_with_ec_conn_output_copies_shared_empty_output(): |
| 317 | """The shared empty output is copied, never written to.""" |
| 318 | ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) |
| 319 | |
| 320 | result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output) |
| 321 | |
| 322 | assert result is not EMPTY_MODEL_RUNNER_OUTPUT |
| 323 | assert result.ec_connector_output is ec_output |
| 324 | assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None |