| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Routing-math tests for batch-sharded sampling. |
| 4 | |
| 5 | These simulate every TP rank in-process (no torch.distributed): the local |
| 6 | sub-batch, the all-to-all send/recv layout, and the padded result gather must |
| 7 | compose to a bit-exact round trip of the replicated computation. |
| 8 | """ |
| 9 | |
| 10 | from types import SimpleNamespace |
| 11 | from unittest import mock |
| 12 | |
| 13 | import numpy as np |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.v1.core.sched.output import GrammarOutput |
| 18 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 19 | from vllm.v1.worker.gpu.sample import batch_shard |
| 20 | from vllm.v1.worker.gpu.sample.batch_shard import ( |
| 21 | BatchSharder, |
| 22 | _shard_grammar_output, |
| 23 | ) |
| 24 | |
| 25 | DEVICE = "cuda" |
| 26 | |
| 27 | # The owner-sorted layout is expanded by a Triton kernel, so the routing math |
| 28 | # can only be exercised on device. |
| 29 | requires_cuda = pytest.mark.skipif( |
| 30 | not torch.cuda.is_available(), reason="requires CUDA" |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | def _np(t: torch.Tensor) -> np.ndarray: |
| 35 | return t.cpu().numpy() |
| 36 | |
| 37 | |
| 38 | def _make_batch( |
| 39 | rng: np.random.Generator, |
| 40 | num_reqs: int, |
| 41 | max_num_reqs: int, |
| 42 | max_spec: int, |
| 43 | slot_pool: np.ndarray | None = None, |
| 44 | ) -> InputBatch: |
| 45 | """Consistent batch with random unique request slots, 1 + k_i logits |
| 46 | rows per request at the tail of each query segment.""" |
| 47 | if slot_pool is None: |
| 48 | slot_pool = np.arange(max_num_reqs) |
| 49 | slots = rng.choice(slot_pool, size=num_reqs, replace=False) |
| 50 | slots = slots.astype(np.int32) |
| 51 | k = rng.integers(0, max_spec + 1, size=num_reqs) |
| 52 | num_logits_per_req = (1 + k).astype(np.int32) |
| 53 | cu = np.zeros(num_reqs + 1, dtype=np.int32) |
| 54 | np.cumsum(num_logits_per_req, out=cu[1:]) |
| 55 | |
| 56 | query_lens = num_logits_per_req + rng.integers(0, 4, size=num_reqs).astype(np.int32) |
| 57 | qsl = np.zeros(num_reqs + 1, dtype=np.int32) |
| 58 | np.cumsum(query_lens, out=qsl[1:]) |
| 59 | num_tokens = int(qsl[-1]) |
| 60 | |
| 61 | logits_indices_np = np.concatenate( |
| 62 | [ |
| 63 | np.arange(qsl[i + 1] - num_logits_per_req[i], qsl[i + 1]) |
| 64 | for i in range(num_reqs) |
| 65 | ] |
| 66 | ).astype(np.int64) |
| 67 | expanded_idx_np = np.repeat(slots, num_logits_per_req) |
| 68 | expanded_pos_np = np.concatenate( |
| 69 | [np.arange(n, dtype=np.int32) for n in num_logits_per_req] |
| 70 | ) |
| 71 | seq_lens_np = (query_lens + rng.integers(0, 100, size=num_reqs)).astype(np.int32) |
| 72 | |
| 73 | def dev(a: np.ndarray) -> torch.Tensor: |
| 74 | return torch.from_numpy(a).to(DEVICE) |
| 75 | |
| 76 | return InputBatch( |
| 77 | req_ids=[f"req-{i}" for i in range(num_reqs)], |
| 78 | num_reqs=num_reqs, |
| 79 | num_reqs_after_padding=num_reqs, |
| 80 | idx_mapping=dev(slots), |
| 81 | idx_mapping_np=slots, |
| 82 | expanded_idx_mapping=dev(expanded_idx_np), |
| 83 | expanded_local_pos=dev(expanded_pos_np), |
| 84 | num_scheduled_tokens=query_lens, |
| 85 | num_tokens=num_tokens, |
| 86 | num_tokens_after_padding=num_tokens, |
| 87 | num_draft_tokens=int(k.sum()), |
| 88 | num_draft_tokens_per_req=k.astype(np.int32), |
| 89 | query_start_loc=dev(qsl), |
| 90 | query_start_loc_np=qsl, |
| 91 | seq_lens=dev(seq_lens_np), |
| 92 | seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_np.copy()), |
| 93 | dcp_local_seq_lens=None, |
| 94 | num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32), |
| 95 | prefill_len_np=np.zeros(num_reqs, dtype=np.int32), |
| 96 | num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), |
| 97 | is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), |
| 98 | has_prefill=False, |
| 99 | decode_graph_eligible=True, |
| 100 | input_ids=torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE), |
| 101 | positions=torch.arange(num_tokens, dtype=torch.int64, device=DEVICE), |
| 102 | is_padding=torch.zeros(num_tokens, dtype=torch.bool, device=DEVICE), |
| 103 | logits_indices=dev(logits_indices_np), |
| 104 | cu_num_logits=dev(cu), |
| 105 | cu_num_logits_np=cu, |
| 106 | has_structured_output_reqs=False, |
| 107 | prompt_lens=None, |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | def _shard_all_ranks(batch: InputBatch, max_num_reqs: int, tp_size: int): |
| 112 | """Simulate every TP rank by faking get_tp_group() per BatchSharder.""" |
| 113 | results = [] |
| 114 | for rank in range(tp_size): |
| 115 | fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size) |
| 116 | with mock.patch.object(batch_shard, "get_tp_group", return_value=fake_group): |
| 117 | sharder = BatchSharder( |
| 118 | max_num_reqs=max_num_reqs, |
| 119 | max_num_logits_per_req=8, |
| 120 | device=torch.device(DEVICE), |
| 121 | ) |
| 122 | local_batch, sorted_logits_indices, _, metadata = sharder.shard_sampler_inputs( |
| 123 | batch, None |
| 124 | ) |
| 125 | results.append((local_batch, sorted_logits_indices, metadata)) |
| 126 | return results |
| 127 | |
| 128 | |
| 129 | def _owned_batch_indices(local_batch: InputBatch) -> np.ndarray: |
| 130 | return np.array([int(r.split("-")[1]) for r in local_batch.req_ids], dtype=np.int64) |
| 131 | |
| 132 | |
| 133 | def _owned_rows(cu: np.ndarray, owned: np.ndarray) -> np.ndarray: |
| 134 | if len(owned) == 0: |
| 135 | return np.array([], dtype=np.int64) |
| 136 | return np.concatenate([np.arange(cu[r], cu[r + 1]) for r in owned]) |
| 137 | |
| 138 | |
| 139 | @requires_cuda |
| 140 | @pytest.mark.parametrize("tp_size", [2, 4, 8]) |
| 141 | @pytest.mark.parametrize("seed", [0, 1]) |
| 142 | def test_local_batch_partition(tp_size: int, seed: int): |
| 143 | """Local sub-batches partition the batch by owner, preserve batch order |
| 144 | within an owner, and carry exactly the owned requests' entries of every |
| 145 | per-request and per-row array.""" |
| 146 | rng = np.random.default_rng(seed) |
| 147 | max_num_reqs = 64 |
| 148 | batch = _make_batch(rng, num_reqs=23, max_num_reqs=max_num_reqs, max_spec=3) |
| 149 | num_logits = int(batch.cu_num_logits_np[-1]) |
| 150 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 151 | cu = batch.cu_num_logits_np |
| 152 | |
| 153 | metadata0 = results[0][2] |
| 154 | assert sum(metadata0.num_logits_per_rank) == num_logits |
| 155 | # Gather width = max logits rows of any request (bounds num_sampled). |
| 156 | assert metadata0.max_num_logits_per_req == int(np.diff(cu).max()) |
| 157 | |
| 158 | all_owned: list[int] = [] |
| 159 | for rank, (local, _, metadata) in enumerate(results): |
| 160 | owned = _owned_batch_indices(local) |
| 161 | assert local.num_reqs == len(owned) == metadata.num_local_reqs |
| 162 | assert (batch.idx_mapping_np[owned] % tp_size == rank).all() |
| 163 | # Stable sort: batch order is preserved within an owner. |
| 164 | assert (np.diff(owned) > 0).all() |
| 165 | all_owned.extend(owned.tolist()) |
| 166 | |
| 167 | rows = _owned_rows(cu, owned) |
| 168 | assert ( |
| 169 | metadata.num_local_logits |
| 170 | == len(rows) |
| 171 | == metadata0.num_logits_per_rank[rank] |
| 172 | ) |
| 173 | |
| 174 | expected_idx = batch.idx_mapping_np[owned] |
| 175 | assert (local.idx_mapping_np == expected_idx).all() |
| 176 | assert _np(local.idx_mapping).tolist() == expected_idx.tolist() |
| 177 | assert (np.diff(local.cu_num_logits_np) == np.diff(cu)[owned]).all() |
| 178 | assert _np(local.cu_num_logits).tolist() == local.cu_num_logits_np.tolist() |
| 179 | assert ( |
| 180 | _np(local.expanded_idx_mapping) == _np(batch.expanded_idx_mapping)[rows] |
| 181 | ).all() |
| 182 | assert ( |
| 183 | _np(local.expanded_local_pos) == _np(batch.expanded_local_pos)[rows] |
| 184 | ).all() |
| 185 | assert (_np(local.logits_indices) == _np(batch.logits_indices)[rows]).all() |
| 186 | assert (_np(local.seq_lens) == _np(batch.seq_lens)[owned]).all() |
| 187 | # Indexed by batch position, like seq_lens: a consumer holding local row |
| 188 | # indices reads another request's length if this is not localised too. |
| 189 | assert ( |
| 190 | _np(local.seq_lens_cpu_upper_bound) |
| 191 | == _np(batch.seq_lens_cpu_upper_bound)[owned] |
| 192 | ).all() |
| 193 | assert batch.num_draft_tokens_per_req is not None |
| 194 | assert local.num_draft_tokens_per_req is not None |
| 195 | assert ( |
| 196 | local.num_draft_tokens_per_req == batch.num_draft_tokens_per_req[owned] |
| 197 | ).all() |
| 198 | assert local.num_draft_tokens == int( |
| 199 | batch.num_draft_tokens_per_req[owned].sum() |
| 200 | ) |
| 201 | # Triton specializes kernels on pointer alignment: every packed-buffer |
| 202 | # view fed to kernels must stay 16-byte aligned across batch shapes, |
| 203 | # or serving JITs new kernel variants mid-collective. |
| 204 | for t in ( |
| 205 | local.idx_mapping, |
| 206 | local.expanded_idx_mapping, |
| 207 | local.expanded_local_pos, |
| 208 | local.logits_indices, |
| 209 | local.cu_num_logits, |
| 210 | metadata.gathered_src_indices, |
| 211 | ): |
| 212 | assert t.data_ptr() % 16 == 0 |
| 213 | assert sorted(all_owned) == list(range(batch.num_reqs)) |
| 214 | |
| 215 | # sorted_logits_indices is identical on all ranks: the owner-sorted logits row |
| 216 | # positions, i.e. logits_indices permuted rank-major. |
| 217 | sorted_logits_indices = _np(results[0][1]) |
| 218 | for _, gi, _ in results: |
| 219 | assert (_np(gi) == sorted_logits_indices).all() |
| 220 | rows_all = np.concatenate( |
| 221 | [_owned_rows(cu, _owned_batch_indices(local)) for local, _, _ in results] |
| 222 | ) |
| 223 | assert ( |
| 224 | sorted_logits_indices.tolist() == _np(batch.logits_indices)[rows_all].tolist() |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | @requires_cuda |
| 229 | @pytest.mark.parametrize("tp_size", [2, 4]) |
| 230 | def test_all_to_all_reassembly_simulated(tp_size: int): |
| 231 | """Vocab shards computed from owner-sorted hidden states and routed |
| 232 | through the simulated all-to-all reassemble into exactly the full-vocab |
| 233 | logits of each rank's owned rows.""" |
| 234 | rng = np.random.default_rng(7) |
| 235 | max_num_reqs = 32 |
| 236 | batch = _make_batch(rng, num_reqs=11, max_num_reqs=max_num_reqs, max_spec=4) |
| 237 | vocab = 64 * tp_size |
| 238 | shard_w = vocab // tp_size |
| 239 | token_logits = torch.randn(batch.num_tokens, vocab, device=DEVICE) |
| 240 | |
| 241 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 242 | sorted_logits_indices = results[0][1] |
| 243 | num_logits_per_rank = results[0][2].num_logits_per_rank |
| 244 | bounds = np.zeros(tp_size + 1, dtype=np.int64) |
| 245 | np.cumsum(num_logits_per_rank, out=bounds[1:]) |
| 246 | |
| 247 | # send_r = rank r's vocab shard of the owner-sorted rows, with no further |
| 248 | # permute; chunk d of send_r goes to rank d. |
| 249 | sends = [ |
| 250 | token_logits[sorted_logits_indices][:, r * shard_w : (r + 1) * shard_w] |
| 251 | for r in range(tp_size) |
| 252 | ] |
| 253 | for rank, (local, _, metadata) in enumerate(results): |
| 254 | recv = torch.cat( |
| 255 | [sends[src][bounds[rank] : bounds[rank + 1]] for src in range(tp_size)] |
| 256 | ) |
| 257 | # Reassemble source-rank-major chunks into full-vocab rows, mirroring |
| 258 | # all_to_all_logits (source-rank order is vocab-shard order). |
| 259 | reassembled = ( |
| 260 | recv.view(tp_size, metadata.num_local_logits, shard_w) |
| 261 | .permute(1, 0, 2) |
| 262 | .reshape(metadata.num_local_logits, vocab) |
| 263 | ) |
| 264 | expected = token_logits[local.logits_indices] |
| 265 | assert torch.equal(reassembled, expected) |
| 266 | |
| 267 | |
| 268 | @requires_cuda |
| 269 | def test_shard_no_spec(): |
| 270 | """Batches without draft tokens (one logits row per request) degenerate |
| 271 | cleanly: rows mirror requests.""" |
| 272 | rng = np.random.default_rng(11) |
| 273 | tp_size = 4 |
| 274 | max_num_reqs = 32 |
| 275 | batch = _make_batch(rng, num_reqs=10, max_num_reqs=max_num_reqs, max_spec=0) |
| 276 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 277 | |
| 278 | assert results[0][2].max_num_logits_per_req == 1 |
| 279 | for local, _, metadata in results: |
| 280 | assert metadata.num_local_logits == metadata.num_local_reqs == local.num_reqs |
| 281 | owned = _owned_batch_indices(local) |
| 282 | assert _np(local.logits_indices).tolist() == ( |
| 283 | _np(batch.logits_indices)[owned].tolist() |
| 284 | ) |
| 285 | |
| 286 | |
| 287 | @pytest.mark.parametrize("tp_size", [2, 4]) |
| 288 | def test_shard_grammar_output(tp_size: int): |
| 289 | """The grammar filter keeps exactly the owned requests' bitmask rows, in |
| 290 | grammar order, and the per-rank pieces partition the global bitmask.""" |
| 291 | rng = np.random.default_rng(5) |
| 292 | max_num_reqs = 32 |
| 293 | num_reqs = 12 |
| 294 | slots = rng.choice(np.arange(max_num_reqs), size=num_reqs, replace=False) |
| 295 | k = rng.integers(0, 4, size=num_reqs) |
| 296 | cu = np.zeros(num_reqs + 1, dtype=np.int32) |
| 297 | np.cumsum(1 + k, out=cu[1:]) |
| 298 | owner = slots % tp_size |
| 299 | |
| 300 | # Grammar requests: every other request, in a scrambled order. |
| 301 | req_ids = [f"req-{i}" for i in range(num_reqs)] |
| 302 | grammar_idx = list(rng.permutation(np.arange(0, num_reqs, 2))) |
| 303 | grammar_req_ids = [req_ids[i] for i in grammar_idx] |
| 304 | num_grammar_rows = sum(int(cu[i + 1] - cu[i]) for i in grammar_idx) |
| 305 | bitmask = rng.integers(0, 2**31, size=(num_grammar_rows, 4), dtype=np.int32) |
| 306 | input_batch = SimpleNamespace(req_ids=req_ids, cu_num_logits_np=cu) |
| 307 | |
| 308 | # Global row offset of each grammar request within the bitmask. |
| 309 | row_offsets = {} |
| 310 | cursor = 0 |
| 311 | for i in grammar_idx: |
| 312 | row_offsets[i] = cursor |
| 313 | cursor += int(cu[i + 1] - cu[i]) |
| 314 | |
| 315 | grammar_output = GrammarOutput( |
| 316 | structured_output_request_ids=grammar_req_ids, |
| 317 | grammar_bitmask=bitmask, |
| 318 | ) |
| 319 | |
| 320 | kept_total = 0 |
| 321 | for rank in range(tp_size): |
| 322 | local_req_ids = [req_ids[i] for i in range(num_reqs) if owner[i] == rank] |
| 323 | local = _shard_grammar_output(grammar_output, input_batch, local_req_ids) |
| 324 | expected_idx = [i for i in grammar_idx if owner[i] == rank] |
| 325 | if not expected_idx: |
| 326 | assert local is None |
| 327 | continue |
| 328 | assert local is not None |
| 329 | assert local.structured_output_request_ids == [req_ids[i] for i in expected_idx] |
| 330 | expected_rows = np.concatenate( |
| 331 | [ |
| 332 | np.arange(row_offsets[i], row_offsets[i] + cu[i + 1] - cu[i]) |
| 333 | for i in expected_idx |
| 334 | ] |
| 335 | ) |
| 336 | assert (local.grammar_bitmask == bitmask[expected_rows.astype(int)]).all() |
| 337 | kept_total += local.grammar_bitmask.shape[0] |
| 338 | assert kept_total == num_grammar_rows |
| 339 | |
| 340 | |
| 341 | @requires_cuda |
| 342 | @pytest.mark.parametrize("tp_size", [2, 4]) |
| 343 | def test_gather_src_indices_round_trip(tp_size: int): |
| 344 | """Padded per-rank result blocks indexed by gathered_src_indices restore |
| 345 | batch order, including when some rank owns zero requests.""" |
| 346 | rng = np.random.default_rng(3) |
| 347 | max_num_reqs = 40 |
| 348 | all_slots = np.arange(max_num_reqs) |
| 349 | batch = _make_batch( |
| 350 | rng, |
| 351 | num_reqs=17, |
| 352 | max_num_reqs=max_num_reqs, |
| 353 | max_spec=2, |
| 354 | slot_pool=all_slots[all_slots % tp_size != 0], # rank 0 owns nothing |
| 355 | ) |
| 356 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 357 | metadata0 = results[0][2] |
| 358 | assert metadata0.num_local_reqs == 0 |
| 359 | |
| 360 | # Value of request j is j; each rank packs its owned requests' values into |
| 361 | # a zero-padded (max_num_reqs_per_rank, 1) block, gathered rank-major. |
| 362 | pad = metadata0.max_num_reqs_per_rank |
| 363 | blocks = [] |
| 364 | for local, _, _ in results: |
| 365 | owned = _owned_batch_indices(local) |
| 366 | block = torch.zeros(pad, 1, dtype=torch.int64, device=DEVICE) |
| 367 | block[: len(owned), 0] = torch.from_numpy(owned).to(DEVICE) |
| 368 | blocks.append(block) |
| 369 | gathered = torch.cat(blocks) |
| 370 | |
| 371 | restored = gathered[metadata0.gathered_src_indices][:, 0] |
| 372 | assert restored.tolist() == list(range(batch.num_reqs)) |
| 373 | |
| 374 | |
| 375 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 376 | @pytest.mark.parametrize("tp_size", [2, 4]) |
| 377 | def test_gather_sampler_output_kernels(tp_size: int): |
| 378 | """The fused pack/unpack kernels round-trip per-rank sampler results |
| 379 | into batch order: token ids zero-filled past the local width, counts |
| 380 | cast back to int32, identical on every rank.""" |
| 381 | from dataclasses import replace as dc_replace |
| 382 | |
| 383 | from vllm.v1.worker.gpu.sample.output import SamplerOutput |
| 384 | |
| 385 | rng = np.random.default_rng(9) |
| 386 | max_num_reqs = 32 |
| 387 | batch = _make_batch(rng, num_reqs=13, max_num_reqs=max_num_reqs, max_spec=3) |
| 388 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 389 | device = torch.device("cuda") |
| 390 | num_reqs = batch.num_reqs |
| 391 | width = results[0][2].max_num_logits_per_req |
| 392 | |
| 393 | # Distinctive per-rank local outputs. Even ranks mimic the regular |
| 394 | # sampler (one column, exercising zero-fill); odd ranks mimic the |
| 395 | # rejection sampler's over-allocation (exercising truncation). |
| 396 | local_outputs: list[SamplerOutput | None] = [] |
| 397 | expected_ids = torch.zeros(num_reqs, width, dtype=torch.int64) |
| 398 | expected_sampled = torch.zeros(num_reqs, dtype=torch.int32) |
| 399 | expected_rejected = torch.zeros(num_reqs, dtype=torch.int32) |
| 400 | for rank, (local, _, metadata) in enumerate(results): |
| 401 | if metadata.num_local_reqs == 0: |
| 402 | local_outputs.append(None) |
| 403 | continue |
| 404 | owned = torch.from_numpy(_owned_batch_indices(local)) |
| 405 | src_width = 1 if rank % 2 == 0 else width + 3 |
| 406 | ids = owned[:, None] * 1000 + torch.arange(src_width)[None, :] |
| 407 | num_sampled = (owned % width + 1).to(torch.int32) |
| 408 | num_rejected = (owned % 3).to(torch.int32) |
| 409 | local_outputs.append( |
| 410 | SamplerOutput( |
| 411 | sampled_token_ids=ids.to(device), |
| 412 | logprobs_tensors=None, |
| 413 | num_nans=None, |
| 414 | num_sampled=num_sampled.to(device), |
| 415 | num_rejected=num_rejected.to(device), |
| 416 | ) |
| 417 | ) |
| 418 | copy_width = min(src_width, width) |
| 419 | expected_ids[owned, :copy_width] = ids[:, :copy_width] |
| 420 | expected_sampled[owned] = num_sampled |
| 421 | expected_rejected[owned] = num_rejected |
| 422 | |
| 423 | # Two passes over a mocked all-gather: capture each rank's packed send |
| 424 | # block, then replay with the concatenated blocks as the gathered result. |
| 425 | packed_blocks: dict[int, torch.Tensor] = {} |
| 426 | gathered_full: list[torch.Tensor] = [] |
| 427 | |
| 428 | def run(rank: int) -> SamplerOutput: |
| 429 | metadata = dc_replace( |
| 430 | results[rank][2], |
| 431 | gathered_src_indices=results[rank][2].gathered_src_indices.to(device), |
| 432 | ) |
| 433 | |
| 434 | def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: |
| 435 | packed_blocks[rank] = x.clone() |
| 436 | if gathered_full: |
| 437 | return gathered_full[0] |
| 438 | return torch.zeros( |
| 439 | tp_size * x.shape[0], x.shape[1], dtype=x.dtype, device=x.device |
| 440 | ) |
| 441 | |
| 442 | with mock.patch.object( |
| 443 | batch_shard, "tensor_model_parallel_all_gather", fake_all_gather |
| 444 | ): |
| 445 | return batch_shard.gather_sampler_output( |
| 446 | local_outputs[rank], |
| 447 | metadata, |
| 448 | device, |
| 449 | global_batch=batch, |
| 450 | local_batch=results[rank][0], |
| 451 | ) |
| 452 | |
| 453 | for rank in range(tp_size): |
| 454 | run(rank) |
| 455 | # Zero the uninitialized padding rows so the comparison is deterministic; |
| 456 | # the real path never reads them (gathered_src_indices skips padding). |
| 457 | for rank, (local, _, metadata) in enumerate(results): |
| 458 | packed_blocks[rank][metadata.num_local_reqs :] = 0 |
| 459 | gathered_full.append(torch.cat([packed_blocks[r] for r in range(tp_size)])) |
| 460 | |
| 461 | for rank in range(tp_size): |
| 462 | out = run(rank) |
| 463 | assert out.sampled_token_ids.dtype == torch.int64 |
| 464 | assert out.num_sampled is not None and out.num_sampled.dtype == torch.int32 |
| 465 | assert out.num_rejected is not None |
| 466 | assert torch.equal(out.sampled_token_ids.cpu(), expected_ids) |
| 467 | assert torch.equal(out.num_sampled.cpu(), expected_sampled) |
| 468 | assert torch.equal(out.num_rejected.cpu(), expected_rejected.to(torch.int32)) |
| 469 | |
| 470 | |
| 471 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 472 | def test_gather_sampler_output_logprobs_and_nans(): |
| 473 | """num_nans is reduced per request and gathered; LogprobsTensors are |
| 474 | packed into fixed per-request blocks, width-aligned across ranks (pad |
| 475 | and truncate paths), and restored to global row order.""" |
| 476 | from dataclasses import replace as dc_replace |
| 477 | |
| 478 | from vllm.v1.outputs import LogprobsTensors |
| 479 | from vllm.v1.worker.gpu.sample.output import SamplerOutput |
| 480 | |
| 481 | tp_size = 2 |
| 482 | rng = np.random.default_rng(21) |
| 483 | max_num_reqs = 32 |
| 484 | batch = _make_batch(rng, num_reqs=6, max_num_reqs=max_num_reqs, max_spec=3) |
| 485 | results = _shard_all_ranks(batch, max_num_reqs, tp_size) |
| 486 | device = torch.device("cuda") |
| 487 | num_reqs = batch.num_reqs |
| 488 | cu = batch.cu_num_logits_np |
| 489 | num_logits = int(cu[-1]) |
| 490 | width = results[0][2].max_num_logits_per_req |
| 491 | num_cols = 3 |
| 492 | |
| 493 | def cuda_fields(b: InputBatch) -> InputBatch: |
| 494 | return dc_replace( |
| 495 | b, |
| 496 | cu_num_logits=b.cu_num_logits.to(device), |
| 497 | expanded_local_pos=b.expanded_local_pos.to(device), |
| 498 | ) |
| 499 | |
| 500 | global_cuda = cuda_fields(batch) |
| 501 | |
| 502 | local_outputs: list[SamplerOutput | None] = [] |
| 503 | local_batches: list[InputBatch] = [] |
| 504 | expected_ids = torch.zeros(num_logits, num_cols, dtype=torch.int64) |
| 505 | expected_logprobs = torch.full( |
| 506 | (num_logits, num_cols), float("-inf"), dtype=torch.float32 |
| 507 | ) |
| 508 | expected_ranks = torch.zeros(num_logits, dtype=torch.int64) |
| 509 | expected_nans = torch.zeros(num_reqs, dtype=torch.int32) |
| 510 | for rank, (local, _, metadata) in enumerate(results): |
| 511 | local_batches.append(cuda_fields(local)) |
| 512 | if metadata.num_local_reqs == 0: |
| 513 | local_outputs.append(None) |
| 514 | continue |
| 515 | owned = _owned_batch_indices(local) |
| 516 | rows = torch.from_numpy(_owned_rows(cu, owned)) |
| 517 | # Rank 0 sends fewer columns (pad path), rank 1 more (truncate path). |
| 518 | local_cols = num_cols - 1 if rank == 0 else num_cols + 2 |
| 519 | ids = rows[:, None] * 10 + torch.arange(local_cols)[None, :] |
| 520 | logprobs = ids.float() * 0.5 |
| 521 | ranks_t = rows + 7 |
| 522 | nans = (rows % 4).to(torch.int32) |
| 523 | local_outputs.append( |
| 524 | SamplerOutput( |
| 525 | sampled_token_ids=torch.zeros( |
| 526 | metadata.num_local_reqs, width, dtype=torch.int64, device=device |
| 527 | ), |
| 528 | logprobs_tensors=LogprobsTensors( |
| 529 | logprob_token_ids=ids.to(device), |
| 530 | logprobs=logprobs.to(device), |
| 531 | selected_token_ranks=ranks_t.to(device), |
| 532 | cu_num_generated_tokens=None, |
| 533 | ), |
| 534 | num_nans=nans.to(device), |
| 535 | num_sampled=torch.ones( |
| 536 | metadata.num_local_reqs, dtype=torch.int32, device=device |
| 537 | ), |
| 538 | num_rejected=torch.zeros( |
| 539 | metadata.num_local_reqs, dtype=torch.int32, device=device |
| 540 | ), |
| 541 | ) |
| 542 | ) |
| 543 | copy_cols = min(local_cols, num_cols) |
| 544 | expected_ids[rows, :copy_cols] = ids[:, :copy_cols] |
| 545 | expected_logprobs[rows, :copy_cols] = logprobs[:, :copy_cols] |
| 546 | expected_ranks[rows] = ranks_t |
| 547 | for j in owned: |
| 548 | in_req = (rows >= int(cu[j])) & (rows < int(cu[j + 1])) |
| 549 | expected_nans[j] = int(nans[in_req].sum()) |
| 550 | |
| 551 | recorded: dict[int, list[torch.Tensor]] = {} |
| 552 | gathered_full: list[torch.Tensor] = [] |
| 553 | |
| 554 | def run(rank: int) -> SamplerOutput: |
| 555 | metadata = dc_replace( |
| 556 | results[rank][2], |
| 557 | gathered_src_indices=results[rank][2].gathered_src_indices.to(device), |
| 558 | ) |
| 559 | calls = {"i": 0} |
| 560 | |
| 561 | def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor: |
| 562 | i = calls["i"] |
| 563 | calls["i"] += 1 |
| 564 | if gathered_full: |
| 565 | return gathered_full[i] |
| 566 | recorded.setdefault(rank, []).append(x.clone()) |
| 567 | return torch.zeros( |
| 568 | tp_size * x.shape[0], *x.shape[1:], dtype=x.dtype, device=x.device |
| 569 | ) |
| 570 | |
| 571 | with mock.patch.object( |
| 572 | batch_shard, "tensor_model_parallel_all_gather", fake_all_gather |
| 573 | ): |
| 574 | return batch_shard.gather_sampler_output( |
| 575 | local_outputs[rank], |
| 576 | metadata, |
| 577 | device, |
| 578 | global_batch=global_cuda, |
| 579 | local_batch=local_batches[rank], |
| 580 | gather_num_nans=True, |
| 581 | logprobs_dims=(num_cols - 1, 0), |
| 582 | ) |
| 583 | |
| 584 | for rank in range(tp_size): |
| 585 | run(rank) |
| 586 | gathered_full.extend( |
| 587 | torch.cat([recorded[r][i] for r in range(tp_size)]) for i in range(3) |
| 588 | ) |
| 589 | |
| 590 | for rank in range(tp_size): |
| 591 | out = run(rank) |
| 592 | assert out.num_nans is not None |
| 593 | assert torch.equal(out.num_nans.cpu(), expected_nans) |
| 594 | lp = out.logprobs_tensors |
| 595 | assert lp is not None |
| 596 | assert torch.equal(lp.logprob_token_ids.cpu(), expected_ids) |
| 597 | assert torch.equal(lp.logprobs.cpu(), expected_logprobs) |
| 598 | assert torch.equal(lp.selected_token_ranks.cpu(), expected_ranks) |
| 599 | if num_logits != num_reqs: |
| 600 | assert lp.cu_num_generated_tokens == cu.tolist() |
| 601 | else: |
| 602 | assert lp.cu_num_generated_tokens is None |
| 603 | |
| 604 | |
| 605 | def test_finish_requests_frees_slots_in_sorted_order(): |
| 606 | """Request slots must be freed in the same order on every TP rank: |
| 607 | ownership derives from slot indices, and `finished_req_ids` is a set |
| 608 | whose iteration order is per-process hash-randomized.""" |
| 609 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 610 | |
| 611 | removed: list[str] = [] |
| 612 | fake_runner = SimpleNamespace(_remove_request=removed.append, pooling_runner=None) |
| 613 | scheduler_output = SimpleNamespace( |
| 614 | finished_req_ids={"req-b", "req-c", "req-a"}, |
| 615 | preempted_req_ids={"req-e", "req-d"}, |
| 616 | ) |
| 617 | GPUModelRunner.finish_requests(fake_runner, scheduler_output) |
| 618 | assert removed == sorted(removed) |
| 619 | assert set(removed) == {"req-a", "req-b", "req-c", "req-d", "req-e"} |