@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 139 lines 4.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.distributed.kv_transfer.kv_connector.utils import (
7 EngineTransferInfo,
8 TransferTopology,
9 )
10
11 pytestmark = pytest.mark.cpu_test
12
13
14 class _FakeAttentionBackend:
15 pass
16
17
18 def _make_topology(
19 *,
20 tp_rank: int = 1,
21 tp_size: int = 4,
22 total_num_kv_heads: int = 8,
23 ) -> TransferTopology:
24 return TransferTopology(
25 tp_rank=tp_rank,
26 tp_size=tp_size,
27 block_size=16,
28 engine_id="local-engine",
29 is_mla=False,
30 is_mamba=False,
31 total_num_kv_heads=total_num_kv_heads,
32 attn_backends=[_FakeAttentionBackend],
33 )
34
35
36 def test_legacy_register_remote_engine_uses_pp_rank_zero() -> None:
37 topology = _make_topology()
38 info = EngineTransferInfo(
39 remote_tp_size=2,
40 remote_block_len=1024,
41 remote_block_size=16,
42 remote_physical_blocks_per_logical=1,
43 )
44
45 registered = topology.register_remote_engine("remote-engine", info)
46
47 assert registered == info
48 assert registered.remote_pp_rank == 0
49 assert topology.get_engine_info("remote-engine") == info
50 assert topology._engines[("remote-engine", 0)] == info
51 assert topology.target_remote_ranks("remote-engine") == [0]
52
53
54 def test_register_remote_engine_stores_pp_ranks_separately() -> None:
55 topology = _make_topology(tp_rank=0, tp_size=2)
56
57 info_0 = EngineTransferInfo(
58 remote_tp_size=2,
59 remote_block_len=1024,
60 remote_block_size=16,
61 remote_physical_blocks_per_logical=1,
62 remote_pp_rank=0,
63 start_layer=0,
64 end_layer=16,
65 )
66 info_1 = EngineTransferInfo(
67 remote_tp_size=1,
68 remote_block_len=512,
69 remote_block_size=8,
70 remote_physical_blocks_per_logical=2,
71 remote_pp_rank=1,
72 start_layer=16,
73 end_layer=32,
74 )
75
76 registered_0 = topology.register_remote_engine("remote-engine", info_0)
77 registered_1 = topology.register_remote_engine("remote-engine", info_1)
78
79 assert registered_0 == info_0
80 assert registered_1 == info_1
81 assert topology.get_engine_info("remote-engine") == info_0
82 assert topology.get_engine_info("remote-engine", 0) == info_0
83 assert topology.get_engine_info("remote-engine", 1) == info_1
84 assert set(topology._engines) == {
85 ("remote-engine", 0),
86 ("remote-engine", 1),
87 }
88
89
90 def test_helpers_use_requested_pp_rank() -> None:
91 topology = _make_topology(tp_rank=1, tp_size=2, total_num_kv_heads=2)
92 topology.register_remote_engine(
93 "remote-engine",
94 EngineTransferInfo(
95 remote_tp_size=1,
96 remote_block_len=1024,
97 remote_block_size=16,
98 remote_physical_blocks_per_logical=1,
99 remote_pp_rank=0,
100 start_layer=0,
101 end_layer=8,
102 ),
103 )
104 topology.register_remote_engine(
105 "remote-engine",
106 EngineTransferInfo(
107 remote_tp_size=4,
108 remote_block_len=1024,
109 remote_block_size=16,
110 remote_physical_blocks_per_logical=1,
111 remote_pp_rank=1,
112 start_layer=8,
113 end_layer=16,
114 ),
115 )
116
117 assert not topology.is_kv_replicated("remote-engine", 0)
118 assert topology.is_kv_replicated("remote-engine", 1)
119 assert topology.replicates_kv_cache("remote-engine", 1)
120 assert topology.target_remote_ranks("remote-engine", 0) == [0]
121 assert topology.target_remote_ranks("remote-engine", 1) == [2, 3]
122 assert "remote_pp=1" in topology.describe("remote-engine", 1)
123
124
125 def test_engine_info_fields_have_backward_compatible_defaults() -> None:
126 topology = _make_topology()
127 info = EngineTransferInfo(
128 remote_tp_size=2,
129 remote_block_len=1024,
130 remote_block_size=16,
131 remote_physical_blocks_per_logical=1,
132 )
133
134 registered = topology.register_remote_engine("remote-engine", info)
135
136 assert topology.get_engine_info("remote-engine") == registered
137 assert registered.remote_pp_rank == 0
138 assert registered.start_layer == 0
139 assert registered.end_layer == 0