| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import replace |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import CUDAGraphMode, ParallelConfig |
| 11 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens |
| 12 | from vllm.v1.attention.ops import pcp as attention_pcp |
| 13 | from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils |
| 14 | from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module |
| 15 | from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor |
| 16 | from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context |
| 17 | from vllm.v1.worker.gpu.pcp_manager import PCPManager |
| 18 | |
| 19 | |
| 20 | def _copy_to_cpu(value, device=None, dtype=None, out=None): |
| 21 | tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value |
| 22 | if out is not None: |
| 23 | return out.copy_(tensor) |
| 24 | return tensor |
| 25 | |
| 26 | |
| 27 | def _make_config(cudagraph_mode: CUDAGraphMode): |
| 28 | return SimpleNamespace( |
| 29 | parallel_config=SimpleNamespace( |
| 30 | prefill_context_parallel_size=2, |
| 31 | decode_context_parallel_size=1, |
| 32 | pipeline_parallel_size=1, |
| 33 | dcp_comm_backend="ag_rs", |
| 34 | ), |
| 35 | model_config=SimpleNamespace( |
| 36 | use_mla=True, |
| 37 | is_encoder_decoder=False, |
| 38 | hf_text_config=SimpleNamespace(), |
| 39 | ), |
| 40 | lora_config=None, |
| 41 | speculative_config=None, |
| 42 | compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def _make_capture_manager(block_table: torch.Tensor): |
| 47 | block_tables = SimpleNamespace( |
| 48 | input_block_tables=(block_table,), |
| 49 | num_kv_cache_groups=1, |
| 50 | kernel_block_sizes=(2,), |
| 51 | blocks_per_kv_block=(1,), |
| 52 | ) |
| 53 | manager = PCPManager( |
| 54 | pcp_world_size=2, |
| 55 | pcp_rank=0, |
| 56 | device=torch.device("cpu"), |
| 57 | shard_decode_requests=False, |
| 58 | max_num_reqs=4, |
| 59 | max_num_tokens=8, |
| 60 | block_tables=block_tables, |
| 61 | ) |
| 62 | return manager, block_tables |
| 63 | |
| 64 | |
| 65 | @pytest.mark.parametrize( |
| 66 | "cudagraph_mode", |
| 67 | [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], |
| 68 | ) |
| 69 | def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode): |
| 70 | PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False) |
| 71 | |
| 72 | |
| 73 | def test_validate_config_rejects_full_graph_for_prefills(): |
| 74 | with pytest.raises(NotImplementedError, match="decode-only routines"): |
| 75 | PCPManager.validate_config( |
| 76 | _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | @pytest.mark.parametrize("draft_uses_mla", [False, True]) |
| 81 | @pytest.mark.parametrize("dcp_size", [1, 2, 4]) |
| 82 | def test_validate_config_dcp_for_dspark(draft_uses_mla, dcp_size): |
| 83 | config = _make_config(CUDAGraphMode.NONE) |
| 84 | config.parallel_config.prefill_context_parallel_size = 4 |
| 85 | config.parallel_config.decode_context_parallel_size = dcp_size |
| 86 | config.model_config.hf_text_config.index_topk = 1 |
| 87 | config.speculative_config = SimpleNamespace( |
| 88 | use_dspark=lambda: True, |
| 89 | draft_model_config=SimpleNamespace(use_mla=draft_uses_mla), |
| 90 | ) |
| 91 | if draft_uses_mla and dcp_size == 2: |
| 92 | with pytest.raises(NotImplementedError, match="requires DCP=1 or DCP=PCP"): |
| 93 | PCPManager.validate_config(config, supports_mm_inputs=False) |
| 94 | return |
| 95 | PCPManager.validate_config(config, supports_mm_inputs=False) |
| 96 | |
| 97 | |
| 98 | def test_sharded_decode_piecewise_graph_padding(monkeypatch): |
| 99 | manager = PCPManager( |
| 100 | pcp_world_size=2, |
| 101 | pcp_rank=0, |
| 102 | device=torch.device("cpu"), |
| 103 | shard_decode_requests=True, |
| 104 | dcp_world_size=1, |
| 105 | ) |
| 106 | monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) |
| 107 | |
| 108 | segments_by_rank, per_rank_num_tokens = manager._build_batch_layout( |
| 109 | num_scheduled_tokens=np.ones(3, dtype=np.int32), |
| 110 | num_computed_tokens=np.full(3, 16, dtype=np.int32), |
| 111 | is_prefilling=np.zeros(3, dtype=np.bool_), |
| 112 | query_start_loc_np=np.arange(4, dtype=np.int32), |
| 113 | padded_num_tokens=4, |
| 114 | ) |
| 115 | |
| 116 | assert per_rank_num_tokens == [2, 1] |
| 117 | request_indices = [ |
| 118 | [segment.global_batch_req_idx for segment in rank] for rank in segments_by_rank |
| 119 | ] |
| 120 | assert request_indices == [[0, 2], [1]] |
| 121 | assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 4, 1])) |
| 122 | assert torch.equal( |
| 123 | manager._padded_gather_idx, |
| 124 | torch.tensor([0, 2, 0, 0, 1, 0, 0, 0]), |
| 125 | ) |
| 126 | assert torch.equal( |
| 127 | manager._gathered_kv_write_mask, |
| 128 | torch.tensor([True, True, False, False, True, False, False, False]), |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | def test_input_buffers_are_exposed_for_cudagraph_capture(): |
| 133 | manager = PCPManager( |
| 134 | pcp_world_size=2, |
| 135 | pcp_rank=0, |
| 136 | device=torch.device("cpu"), |
| 137 | shard_decode_requests=False, |
| 138 | max_num_reqs=4, |
| 139 | max_num_tokens=8, |
| 140 | ) |
| 141 | |
| 142 | assert manager.input_buffers is manager._input_buffers |
| 143 | assert manager.input_buffers.input_ids.shape == (8,) |
| 144 | assert manager.input_buffers.positions.shape == (8,) |
| 145 | assert manager.input_buffers.is_padding.shape == (8,) |
| 146 | |
| 147 | |
| 148 | @pytest.mark.parametrize( |
| 149 | ("pcp_world_size", "num_scheduled_tokens", "is_prefilling", "expected"), |
| 150 | [ |
| 151 | (2, [8], [True], 4), |
| 152 | (2, [7], [True], 4), |
| 153 | (2, [3], [False], 3), |
| 154 | (2, [3, 8], [False, True], 7), |
| 155 | (4, [2, 9], [False, True], 4), |
| 156 | ], |
| 157 | ) |
| 158 | def test_num_tokens_for_dispatch_uses_largest_pcp_rank( |
| 159 | pcp_world_size, num_scheduled_tokens, is_prefilling, expected |
| 160 | ): |
| 161 | manager = PCPManager( |
| 162 | pcp_world_size=pcp_world_size, |
| 163 | pcp_rank=0, |
| 164 | device=torch.device("cpu"), |
| 165 | shard_decode_requests=True, |
| 166 | ) |
| 167 | |
| 168 | actual = manager.get_num_tokens_for_dispatch( |
| 169 | np.asarray(num_scheduled_tokens, dtype=np.int32), |
| 170 | np.asarray(is_prefilling, dtype=np.bool_), |
| 171 | ) |
| 172 | |
| 173 | assert actual == expected |
| 174 | |
| 175 | |
| 176 | def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): |
| 177 | manager = PCPManager( |
| 178 | pcp_world_size=2, |
| 179 | pcp_rank=0, |
| 180 | device=torch.device("cpu"), |
| 181 | shard_decode_requests=True, |
| 182 | dcp_world_size=1, |
| 183 | ) |
| 184 | monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) |
| 185 | |
| 186 | with pytest.raises(ValueError, match="smaller than the largest rank-local batch"): |
| 187 | manager._build_batch_layout( |
| 188 | num_scheduled_tokens=np.ones(3, dtype=np.int32), |
| 189 | num_computed_tokens=np.full(3, 16, dtype=np.int32), |
| 190 | is_prefilling=np.zeros(3, dtype=np.bool_), |
| 191 | query_start_loc_np=np.arange(4, dtype=np.int32), |
| 192 | padded_num_tokens=1, |
| 193 | ) |
| 194 | |
| 195 | |
| 196 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels") |
| 197 | @pytest.mark.parametrize( |
| 198 | ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"), |
| 199 | [ |
| 200 | (CUDAGraphMode.NONE, 4, 2, 2), |
| 201 | (CUDAGraphMode.PIECEWISE, None, 8, 2), |
| 202 | (CUDAGraphMode.FULL, 4, 8, 4), |
| 203 | ], |
| 204 | ) |
| 205 | def test_partition_padding_is_derived_from_batch_descriptor( |
| 206 | cg_mode, num_reqs, expected_tokens, expected_reqs |
| 207 | ): |
| 208 | """partition_batch derives graph padding from the batch descriptor: no |
| 209 | padding without graphs, token padding only for PIECEWISE, and request |
| 210 | padding on top for FULL.""" |
| 211 | device = torch.device("cuda:0") |
| 212 | manager = PCPManager( |
| 213 | pcp_world_size=2, |
| 214 | pcp_rank=0, |
| 215 | device=device, |
| 216 | shard_decode_requests=False, |
| 217 | max_num_reqs=4, |
| 218 | max_num_tokens=8, |
| 219 | ) |
| 220 | global_batch = _make_global_decode_batch( |
| 221 | [16, 24], InputBuffers(4, 8, device), device |
| 222 | ) |
| 223 | |
| 224 | local_batch = manager.partition_batch( |
| 225 | global_batch, |
| 226 | BatchExecutionDescriptor( |
| 227 | cg_mode=cg_mode, |
| 228 | num_tokens=8, |
| 229 | num_reqs=num_reqs, |
| 230 | ), |
| 231 | ) |
| 232 | |
| 233 | assert local_batch.num_tokens_after_padding == expected_tokens |
| 234 | assert local_batch.num_reqs_after_padding == expected_reqs |
| 235 | |
| 236 | |
| 237 | def test_dummy_draft_does_not_reuse_previous_graph_batch(): |
| 238 | manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) |
| 239 | dummy_batch = InputBatch.make_dummy(1, 4, manager.input_buffers) |
| 240 | manager.draft_prefill_batch = replace(dummy_batch) |
| 241 | |
| 242 | manager.prepare_draft_prefill(dummy_batch, dummy_batch.input_ids) |
| 243 | |
| 244 | assert ( |
| 245 | manager.get_draft_input_buffers(manager.input_buffers) is manager.input_buffers |
| 246 | ) |
| 247 | |
| 248 | |
| 249 | def test_capture_uses_pcp_persistent_inputs(): |
| 250 | manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) |
| 251 | |
| 252 | dummy_batch = InputBatch.make_dummy( |
| 253 | num_reqs=4, |
| 254 | num_tokens=4, |
| 255 | input_buffers=InputBuffers(4, 8, torch.device("cpu")), |
| 256 | max_query_len=1, |
| 257 | ) |
| 258 | dummy_batch.input_ids.copy_(torch.arange(4)) |
| 259 | dummy_batch.positions.fill_(7) |
| 260 | input_batch = manager.prepare_inputs_to_capture(dummy_batch) |
| 261 | |
| 262 | assert input_batch is not dummy_batch |
| 263 | assert input_batch.req_ids == dummy_batch.req_ids |
| 264 | for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"): |
| 265 | actual = getattr(input_batch, name) |
| 266 | torch.testing.assert_close(actual, getattr(dummy_batch, name)) |
| 267 | assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr() |
| 268 | |
| 269 | |
| 270 | def test_dummy_context_updates_pcp_local_block_tables(): |
| 271 | global_block_table = torch.full((4, 4), -1, dtype=torch.int32) |
| 272 | manager, block_tables = _make_capture_manager(global_block_table) |
| 273 | input_batch = InputBatch.make_dummy( |
| 274 | num_reqs=2, |
| 275 | num_tokens=2, |
| 276 | input_buffers=manager.input_buffers, |
| 277 | max_query_len=1, |
| 278 | ) |
| 279 | input_batch = manager.prepare_inputs_to_capture(input_batch) |
| 280 | local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs) |
| 281 | |
| 282 | set_dummy_context( |
| 283 | input_batch, |
| 284 | block_tables, |
| 285 | context_len=3, |
| 286 | num_kv_blocks=16, |
| 287 | max_model_len=16, |
| 288 | input_block_tables=local_block_tables, |
| 289 | ) |
| 290 | |
| 291 | torch.testing.assert_close( |
| 292 | local_block_tables[0][:2, :2], |
| 293 | torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), |
| 294 | ) |
| 295 | assert torch.all(global_block_table == -1) |
| 296 | |
| 297 | |
| 298 | def _rank_request_ids( |
| 299 | manager: PCPManager, |
| 300 | rank: int, |
| 301 | req_ids: list[str], |
| 302 | *, |
| 303 | is_prefilling: np.ndarray | None = None, |
| 304 | ) -> list[str]: |
| 305 | num_reqs = len(req_ids) |
| 306 | if is_prefilling is None: |
| 307 | is_prefilling = np.zeros(num_reqs, dtype=np.bool_) |
| 308 | segments = manager._get_rank_segments( |
| 309 | rank=rank, |
| 310 | num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32), |
| 311 | is_prefilling=is_prefilling, |
| 312 | query_start_loc_np=np.arange(num_reqs + 1, dtype=np.int32), |
| 313 | ) |
| 314 | return [req_ids[segment.global_batch_req_idx] for segment in segments] |
| 315 | |
| 316 | |
| 317 | def test_pcp_only_decode_requests_are_round_robin_balanced_each_step(): |
| 318 | manager = PCPManager( |
| 319 | pcp_world_size=4, |
| 320 | pcp_rank=0, |
| 321 | device=torch.device("cpu"), |
| 322 | shard_decode_requests=True, |
| 323 | dcp_world_size=1, |
| 324 | ) |
| 325 | req_ids = [f"request-{idx}" for idx in range(18)] |
| 326 | |
| 327 | owners: dict[str, int] = {} |
| 328 | for rank in range(manager.pcp_world_size): |
| 329 | for req_id in _rank_request_ids(manager, rank, req_ids): |
| 330 | assert req_id not in owners |
| 331 | owners[req_id] = rank |
| 332 | |
| 333 | assert owners == { |
| 334 | req_id: index % manager.pcp_world_size for index, req_id in enumerate(req_ids) |
| 335 | } |
| 336 | counts = [list(owners.values()).count(rank) for rank in range(4)] |
| 337 | assert max(counts) - min(counts) == 1 |
| 338 | |
| 339 | reordered_req_ids = req_ids[::2] + req_ids[1::2] |
| 340 | reordered_owners = { |
| 341 | req_id: rank |
| 342 | for rank in range(manager.pcp_world_size) |
| 343 | for req_id in _rank_request_ids(manager, rank, reordered_req_ids) |
| 344 | } |
| 345 | assert reordered_owners == { |
| 346 | req_id: index % manager.pcp_world_size |
| 347 | for index, req_id in enumerate(reordered_req_ids) |
| 348 | } |
| 349 | assert reordered_owners != owners |
| 350 | |
| 351 | |
| 352 | def test_decode_requests_remain_replicated_when_dcp_is_enabled(): |
| 353 | manager = PCPManager( |
| 354 | pcp_world_size=2, |
| 355 | pcp_rank=0, |
| 356 | device=torch.device("cpu"), |
| 357 | shard_decode_requests=False, |
| 358 | dcp_world_size=2, |
| 359 | ) |
| 360 | req_ids = ["request-a", "request-b", "request-c"] |
| 361 | |
| 362 | assert _rank_request_ids(manager, 0, req_ids) == req_ids |
| 363 | assert _rank_request_ids(manager, 1, req_ids) == req_ids |
| 364 | |
| 365 | |
| 366 | def test_decode_sharding_allows_ranks_with_no_owned_request(): |
| 367 | manager = PCPManager( |
| 368 | pcp_world_size=4, |
| 369 | pcp_rank=0, |
| 370 | device=torch.device("cpu"), |
| 371 | shard_decode_requests=True, |
| 372 | dcp_world_size=1, |
| 373 | ) |
| 374 | req_ids = ["request-a"] |
| 375 | |
| 376 | assert _rank_request_ids(manager, 0, req_ids) == req_ids |
| 377 | assert _rank_request_ids(manager, 1, req_ids) == [] |
| 378 | assert _rank_request_ids(manager, 2, req_ids) == [] |
| 379 | assert _rank_request_ids(manager, 3, req_ids) == [] |
| 380 | |
| 381 | |
| 382 | def test_prefill_partitioning_is_preserved_with_sharded_decode(): |
| 383 | manager = PCPManager( |
| 384 | pcp_world_size=2, |
| 385 | pcp_rank=0, |
| 386 | device=torch.device("cpu"), |
| 387 | shard_decode_requests=True, |
| 388 | dcp_world_size=1, |
| 389 | ) |
| 390 | segments_by_rank = [ |
| 391 | manager._get_rank_segments( |
| 392 | rank=rank, |
| 393 | num_scheduled_tokens=np.array([8, 1, 0, 1], dtype=np.int32), |
| 394 | is_prefilling=np.array([True, False, False, False]), |
| 395 | query_start_loc_np=np.array([0, 8, 9, 9, 10], dtype=np.int32), |
| 396 | ) |
| 397 | for rank in range(manager.pcp_world_size) |
| 398 | ] |
| 399 | |
| 400 | prefill_tokens = sorted( |
| 401 | token_idx |
| 402 | for segments in segments_by_rank |
| 403 | for segment in segments |
| 404 | if segment.global_batch_req_idx == 0 |
| 405 | for token_idx in range( |
| 406 | segment.global_batch_slice.start, segment.global_batch_slice.stop |
| 407 | ) |
| 408 | ) |
| 409 | decode_owners = { |
| 410 | segment.global_batch_req_idx: rank |
| 411 | for rank, segments in enumerate(segments_by_rank) |
| 412 | for segment in segments |
| 413 | if segment.global_batch_req_idx in (1, 3) |
| 414 | } |
| 415 | |
| 416 | assert prefill_tokens == list(range(8)) |
| 417 | assert decode_owners == {1: 0, 3: 1} |
| 418 | |
| 419 | |
| 420 | def test_sharded_decode_layout_selects_owner_kv_for_replication(monkeypatch): |
| 421 | manager = PCPManager( |
| 422 | pcp_world_size=2, |
| 423 | pcp_rank=0, |
| 424 | device=torch.device("cpu"), |
| 425 | shard_decode_requests=True, |
| 426 | dcp_world_size=1, |
| 427 | ) |
| 428 | monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) |
| 429 | manager._build_batch_layout( |
| 430 | num_scheduled_tokens=np.array([1, 1, 1], dtype=np.int32), |
| 431 | num_computed_tokens=np.array([16, 16, 16], dtype=np.int32), |
| 432 | is_prefilling=np.array([False, False, False]), |
| 433 | query_start_loc_np=np.array([0, 1, 2, 3], dtype=np.int32), |
| 434 | ) |
| 435 | |
| 436 | gathered_slot_mapping = manager._convert_to_gathered_slot_mappings( |
| 437 | torch.tensor([[123, 456, 789]], dtype=torch.int64) |
| 438 | ) |
| 439 | assert torch.equal( |
| 440 | gathered_slot_mapping, |
| 441 | torch.tensor([[123, 789, 456, PAD_SLOT_ID]], dtype=torch.int64), |
| 442 | ) |
| 443 | assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 2, 1])) |
| 444 | |
| 445 | class FakePCPGroup: |
| 446 | world_size = 2 |
| 447 | |
| 448 | def all_gather(self, tensor: torch.Tensor, dim: int = 0) -> torch.Tensor: |
| 449 | assert dim == 0 |
| 450 | assert tensor.shape == (2, 1) |
| 451 | return torch.tensor([[11.0], [33.0], [22.0], [0.0]]) |
| 452 | |
| 453 | monkeypatch.setattr(attention_pcp, "get_pcp_group", FakePCPGroup) |
| 454 | (gathered_kv,), cache_slot_mapping = attention_pcp._gather_prefill_cache_inputs( |
| 455 | (torch.tensor([[11.0], [33.0]]),), |
| 456 | gathered_slot_mapping[0], |
| 457 | num_decode_tokens=0, |
| 458 | shard_decode_requests=True, |
| 459 | ) |
| 460 | |
| 461 | assert torch.equal(gathered_kv, torch.tensor([[11.0], [33.0], [22.0], [0.0]])) |
| 462 | assert torch.equal(cache_slot_mapping, torch.tensor([123, 789, 456, PAD_SLOT_ID])) |
| 463 | |
| 464 | |
| 465 | def test_replicated_draft_cache_inputs_keep_unexpanded_slot_mapping(monkeypatch): |
| 466 | def unexpected_gather(): |
| 467 | pytest.fail("replicated draft decode must not gather PCP cache inputs") |
| 468 | |
| 469 | monkeypatch.setattr(attention_pcp, "get_pcp_group", unexpected_gather) |
| 470 | kv = torch.tensor([[11.0], [22.0], [0.0]]) |
| 471 | slot_mapping = torch.tensor([123, 456]) |
| 472 | |
| 473 | (cache_kv,), cache_slot_mapping = attention_pcp._gather_prefill_cache_inputs( |
| 474 | (kv,), slot_mapping, num_decode_tokens=2, shard_decode_requests=True |
| 475 | ) |
| 476 | |
| 477 | assert torch.equal(cache_kv, kv[:2]) |
| 478 | assert torch.equal(cache_slot_mapping, slot_mapping) |
| 479 | |
| 480 | |
| 481 | def _rank_rows( |
| 482 | pcp_rank: int, |
| 483 | pcp_world_size: int, |
| 484 | num_scheduled_tokens: np.ndarray, |
| 485 | num_computed_tokens: np.ndarray, |
| 486 | is_prefilling: np.ndarray, |
| 487 | ) -> tuple[np.ndarray, np.ndarray]: |
| 488 | """One PCP+DCP rank's rows as (global request, extent).""" |
| 489 | manager = PCPManager( |
| 490 | pcp_world_size=pcp_world_size, |
| 491 | pcp_rank=pcp_rank, |
| 492 | device=torch.device("cpu"), |
| 493 | shard_decode_requests=False, |
| 494 | dcp_world_size=2, |
| 495 | ) |
| 496 | query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype( |
| 497 | np.int32 |
| 498 | ) |
| 499 | segments = manager._get_rank_segments( |
| 500 | pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np |
| 501 | ) |
| 502 | rows = np.array([segment.global_batch_req_idx for segment in segments]) |
| 503 | extents = (num_computed_tokens + num_scheduled_tokens)[rows] |
| 504 | return rows, extents |
| 505 | |
| 506 | |
| 507 | @pytest.mark.parametrize( |
| 508 | ("num_scheduled_tokens", "num_computed_tokens"), |
| 509 | [ |
| 510 | ([32, 32], [0, 0]), |
| 511 | ([32, 32, 32, 32], [0, 0, 0, 0]), |
| 512 | ([32, 32], [100, 0]), # one continued, one fresh |
| 513 | ([64, 32, 48], [0, 512, 0]), # ragged lengths and mixed contexts |
| 514 | # A replicated prefill BEHIND a split one. Sorting replicated rows first |
| 515 | # put request 1 ahead of request 0. |
| 516 | ([64, 3], [0, 0]), |
| 517 | ([64, 3, 64], [0, 0, 0]), |
| 518 | ], |
| 519 | ) |
| 520 | @pytest.mark.parametrize("pcp_world_size", [2, 4]) |
| 521 | def test_published_row_order_is_identical_on_every_pcp_rank( |
| 522 | num_scheduled_tokens, num_computed_tokens, pcp_world_size |
| 523 | ): |
| 524 | """Every rank must map its rows to the same global requests, in the same order.""" |
| 525 | num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32) |
| 526 | num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32) |
| 527 | is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_) |
| 528 | |
| 529 | orders = [ |
| 530 | _rank_rows( |
| 531 | rank, |
| 532 | pcp_world_size, |
| 533 | num_scheduled_tokens, |
| 534 | num_computed_tokens, |
| 535 | is_prefilling, |
| 536 | )[0] |
| 537 | for rank in range(pcp_world_size) |
| 538 | ] |
| 539 | for rank, order in enumerate(orders[1:], start=1): |
| 540 | assert np.array_equal(orders[0], order), ( |
| 541 | f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}" |
| 542 | ) |
| 543 | # Grouped by request and ascending: what the indexer plan indexes by. |
| 544 | assert np.all(np.diff(orders[0]) >= 0) |
| 545 | |
| 546 | |
| 547 | def test_published_row_order_puts_every_decode_before_every_prefill(): |
| 548 | """split_decodes_and_prefills takes the FIRST prefilling row as the boundary.""" |
| 549 | # req 0 is a continued prefill, req 1 is a decode. |
| 550 | req_idx, _ = _rank_rows( |
| 551 | pcp_rank=0, |
| 552 | pcp_world_size=2, |
| 553 | num_scheduled_tokens=np.array([32, 1], dtype=np.int32), |
| 554 | num_computed_tokens=np.array([100, 20], dtype=np.int32), |
| 555 | is_prefilling=np.array([True, False], dtype=np.bool_), |
| 556 | ) |
| 557 | # The decode (request 1) must be row 0. |
| 558 | assert req_idx[0] == 1 |
| 559 | |
| 560 | |
| 561 | def test_split_prefill_rows_repeat_the_request_and_its_full_extent(): |
| 562 | """Equal adjacent request indices are what let a backend share a KV region.""" |
| 563 | req_idx, extents = _rank_rows( |
| 564 | pcp_rank=1, |
| 565 | pcp_world_size=2, |
| 566 | num_scheduled_tokens=np.array([64, 3], dtype=np.int32), |
| 567 | num_computed_tokens=np.array([100, 0], dtype=np.int32), |
| 568 | is_prefilling=np.ones(2, dtype=np.bool_), |
| 569 | ) |
| 570 | assert req_idx.tolist() == [0, 0, 1] |
| 571 | assert extents.tolist() == [164, 164, 3] |
| 572 | |
| 573 | |
| 574 | @pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9]) |
| 575 | def test_dcp_replicates_prefills_too_short_to_split(query_len): |
| 576 | """A prefill that cannot fill 2*pcp chunks is replicated, not split.""" |
| 577 | pcp_world_size = 2 |
| 578 | num_scheduled_tokens = np.array([query_len], dtype=np.int32) |
| 579 | is_prefilling = np.ones(1, dtype=np.bool_) |
| 580 | |
| 581 | rows_per_rank = [] |
| 582 | for rank in range(pcp_world_size): |
| 583 | manager = PCPManager( |
| 584 | pcp_world_size=pcp_world_size, |
| 585 | pcp_rank=rank, |
| 586 | device=torch.device("cpu"), |
| 587 | shard_decode_requests=False, |
| 588 | dcp_world_size=2, |
| 589 | ) |
| 590 | assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0] |
| 591 | rows = list( |
| 592 | manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling) |
| 593 | ) |
| 594 | assert rows == [(0, 0, query_len)] |
| 595 | rows_per_rank.append(rows) |
| 596 | assert rows_per_rank[0] == rows_per_rank[1] |
| 597 | |
| 598 | |
| 599 | @pytest.mark.parametrize("pcp_world_size", [2, 4, 8]) |
| 600 | @pytest.mark.parametrize( |
| 601 | "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097] |
| 602 | ) |
| 603 | def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len): |
| 604 | num_scheduled_tokens = np.array([query_len], dtype=np.int32) |
| 605 | is_prefilling = np.ones(1, dtype=np.bool_) |
| 606 | |
| 607 | charged_per_rank = [] |
| 608 | for rank in range(pcp_world_size): |
| 609 | manager = PCPManager( |
| 610 | pcp_world_size=pcp_world_size, |
| 611 | pcp_rank=rank, |
| 612 | device=torch.device("cpu"), |
| 613 | shard_decode_requests=False, |
| 614 | dcp_world_size=pcp_world_size, |
| 615 | ) |
| 616 | chunk_lens = [ |
| 617 | chunk_len |
| 618 | for _, _, chunk_len in manager._iter_rank_chunks( |
| 619 | rank, num_scheduled_tokens, is_prefilling |
| 620 | ) |
| 621 | ] |
| 622 | assert chunk_lens, f"rank {rank} got no rows for {query_len=}" |
| 623 | assert chunk_lens[0] == max(chunk_lens), ( |
| 624 | f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}" |
| 625 | ) |
| 626 | charged_per_rank.append(len(chunk_lens) * chunk_lens[0]) |
| 627 | |
| 628 | assert len(set(charged_per_rank)) == 1, ( |
| 629 | f"ranks would chunk differently for {query_len=}: {charged_per_rank}" |
| 630 | ) |
| 631 | |
| 632 | |
| 633 | def _make_global_decode_batch( |
| 634 | num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device |
| 635 | ) -> InputBatch: |
| 636 | """A replicate-decode global batch as `prepare_inputs` would build it.""" |
| 637 | num_reqs = len(num_computed_tokens) |
| 638 | num_tokens = num_reqs |
| 639 | seq_lens_np = np.asarray(num_computed_tokens, dtype=np.int32) + 1 |
| 640 | |
| 641 | base = InputBatch.make_dummy(num_reqs, num_tokens, buffers) |
| 642 | buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device) |
| 643 | buffers.positions[:num_reqs] = torch.tensor(num_computed_tokens, device=device) |
| 644 | query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) |
| 645 | buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to( |
| 646 | device |
| 647 | ) |
| 648 | |
| 649 | return replace( |
| 650 | base, |
| 651 | req_ids=[f"req_{i}" for i in range(num_reqs)], |
| 652 | num_reqs=num_reqs, |
| 653 | num_reqs_after_padding=num_reqs, |
| 654 | idx_mapping=torch.arange(num_reqs, dtype=torch.int32, device=device), |
| 655 | idx_mapping_np=np.arange(num_reqs, dtype=np.int32), |
| 656 | num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32), |
| 657 | num_tokens=num_tokens, |
| 658 | num_tokens_after_padding=num_tokens, |
| 659 | num_draft_tokens=0, |
| 660 | num_draft_tokens_per_req=np.zeros(num_reqs, dtype=np.int32), |
| 661 | query_start_loc=buffers.query_start_loc[: num_reqs + 1], |
| 662 | query_start_loc_np=query_start_loc_np, |
| 663 | seq_lens=buffers.seq_lens[:num_reqs], |
| 664 | seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_np), |
| 665 | dcp_local_seq_lens=None, |
| 666 | num_computed_tokens_np=np.asarray(num_computed_tokens, dtype=np.int32), |
| 667 | prefill_len_np=np.zeros(num_reqs, dtype=np.int32), |
| 668 | num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), |
| 669 | is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), |
| 670 | input_ids=buffers.input_ids[:num_tokens], |
| 671 | positions=buffers.positions[:num_tokens], |
| 672 | is_padding=buffers.is_padding[:num_tokens], |
| 673 | prompt_lens=None, |
| 674 | ) |
| 675 | |
| 676 | |
| 677 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels") |
| 678 | def test_partition_defers_dcp_metadata_to_post_partition_batch(): |
| 679 | """DCP-local lengths must derive from the partitioned batch, not the |
| 680 | global one: the partition replaces seq_lens, so a pre-partition value is |
| 681 | stale. partition_batch therefore returns None, and the runtime populates |
| 682 | the field afterwards from the PCP-owned buffers. |
| 683 | """ |
| 684 | device = torch.device("cuda:0") |
| 685 | manager = PCPManager( |
| 686 | pcp_world_size=2, |
| 687 | pcp_rank=0, |
| 688 | device=device, |
| 689 | shard_decode_requests=False, |
| 690 | max_num_reqs=4, |
| 691 | max_num_tokens=8, |
| 692 | dcp_world_size=2, |
| 693 | dcp_rank=0, |
| 694 | ) |
| 695 | |
| 696 | global_buffers = InputBuffers(4, 8, device) |
| 697 | global_batch = _make_global_decode_batch([16, 24], global_buffers, device) |
| 698 | # A leftover from an earlier DCP batch must not survive the partition. |
| 699 | global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2] |
| 700 | global_batch.dcp_local_seq_lens.fill_(-1) |
| 701 | |
| 702 | local_batch = manager.partition_batch( |
| 703 | global_batch, |
| 704 | BatchExecutionDescriptor( |
| 705 | cg_mode=CUDAGraphMode.FULL, |
| 706 | num_tokens=4, |
| 707 | num_reqs=4, |
| 708 | ), |
| 709 | ) |
| 710 | |
| 711 | assert local_batch.dcp_local_seq_lens is None |
| 712 | assert local_batch.seq_lens.tolist() == [17, 25, 0, 0] |
| 713 | assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None |
| 714 | expected_dcp_upper_bound = get_dcp_local_seq_lens( |
| 715 | torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 |
| 716 | ) |
| 717 | assert torch.equal( |
| 718 | local_batch.dcp_local_seq_lens_cpu_upper_bound, |
| 719 | torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]), |
| 720 | ) |
| 721 | |
| 722 | # What execute_model does next: derive DCP metadata from the final batch |
| 723 | # on the PCP-owned buffers. |
| 724 | local_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens( |
| 725 | manager.input_buffers.dcp_local_seq_lens, |
| 726 | local_batch.seq_lens, |
| 727 | local_batch.num_reqs, |
| 728 | dcp_size=2, |
| 729 | dcp_rank=0, |
| 730 | cp_interleave=1, |
| 731 | num_reqs_padded=local_batch.num_reqs_after_padding, |
| 732 | ) |
| 733 | expected = get_dcp_local_seq_lens( |
| 734 | torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1 |
| 735 | ) |
| 736 | assert local_batch.dcp_local_seq_lens is not None |
| 737 | assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected) |
| 738 | |
| 739 | |
| 740 | @pytest.mark.parametrize( |
| 741 | ("pcp_world_size", "dcp_world_size", "expected"), |
| 742 | [(1, 1, False), (2, 1, True), (2, 2, False)], |
| 743 | ) |
| 744 | def test_parallel_config_manages_decode_sharding( |
| 745 | pcp_world_size: int, dcp_world_size: int, expected: bool |
| 746 | ): |
| 747 | parallel_config = ParallelConfig( |
| 748 | prefill_context_parallel_size=pcp_world_size, |
| 749 | decode_context_parallel_size=dcp_world_size, |
| 750 | distributed_executor_backend="ray", |
| 751 | ) |
| 752 | |
| 753 | assert parallel_config.pcp_shard_decode_requests is expected |