@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 161 lines 5.69 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from types import SimpleNamespace
4
5 import pytest
6 import torch
7
8 from vllm.config import ParallelConfig
9 from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
10 from vllm.forward_context import DPMetadata
11 from vllm.v1.attention.ops.pcp import maybe_gather_mla_latent_cache_inputs
12
13
14 @pytest.mark.parametrize(
15 "pcp_size,sp_size,enable_ep,expected",
16 [
17 (1, 1, True, [5, 7]),
18 (1, 2, True, [3, 3, 4, 4]),
19 (2, 1, False, [10, 14]),
20 (2, 1, True, [5, 5, 7, 7]),
21 (2, 2, True, [3, 3, 3, 3, 4, 4, 4, 4]),
22 ],
23 )
24 def test_dispatch_sizes_expand_pcp_before_tp(pcp_size, sp_size, enable_ep, expected):
25 config = ParallelConfig(
26 distributed_executor_backend="mp",
27 data_parallel_size=2,
28 prefill_context_parallel_size=pcp_size,
29 tensor_parallel_size=sp_size,
30 enable_expert_parallel=enable_ep,
31 )
32 metadata = DPMetadata.make(config, 5, torch.tensor([5, 7]))
33 with metadata.sp_local_sizes(sp_size, pcp_size=pcp_size, use_ep=enable_ep) as sizes:
34 assert sizes == expected
35 assert metadata.local_sizes is None
36 assert metadata.num_tokens_across_dp_cpu.tolist() == [5, 7]
37
38
39 @pytest.mark.parametrize(
40 "dp_size,pcp_size,tp_size,use_ep,is_sp,expected",
41 [
42 (2, 2, 1, True, False, "ep"),
43 (2, 2, 2, True, True, "ep"),
44 (2, 2, 2, False, False, "dp"),
45 (2, 1, 2, True, False, "dp"),
46 (1, 2, 2, True, False, "pcp"),
47 (2, 2, 2, True, False, None),
48 ],
49 )
50 def test_dispatch_reuses_existing_groups(
51 dp_size, pcp_size, tp_size, use_ep, is_sp, expected, monkeypatch
52 ):
53 groups = {
54 "dp": SimpleNamespace(world_size=dp_size),
55 "pcp": SimpleNamespace(world_size=pcp_size),
56 "ep": object(),
57 }
58 for name, group in groups.items():
59 monkeypatch.setattr(
60 f"vllm.distributed.device_communicators.all2all.get_{name}_group",
61 lambda group=group: group,
62 )
63 manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
64 manager.dp_world_size = dp_size
65 manager.tp_group = SimpleNamespace(world_size=tp_size)
66 manager.use_ep = use_ep
67 if expected is None:
68 with pytest.raises(AssertionError, match="requires sequence-parallel MoE"):
69 manager._get_comm_group(is_sp)
70 else:
71 assert manager._get_comm_group(is_sp) is groups[expected]
72
73
74 @pytest.mark.parametrize("enable_ep", [False, True])
75 def test_ag_rs_dispatch_and_combine_use_dp_pcp_sizes(monkeypatch, enable_ep):
76 calls = []
77 local_tokens = [20, 21] if enable_ep else [20, 21, 22, 23]
78
79 class FakeGroup:
80 world_size = 4 if enable_ep else 2
81 rank_in_group = 2 if enable_ep else 1
82
83 def all_gatherv(self, tensors, dim, sizes):
84 calls.append(("gather", sizes))
85 return [torch.tensor([10, 11, 20, 21, 22, 23]) for _ in tensors]
86
87 def reduce_scatterv(self, tensor, dim, sizes):
88 calls.append(("scatter", sizes))
89 return tensor[2 : 2 + len(local_tokens)]
90
91 config = ParallelConfig(
92 distributed_executor_backend="mp",
93 data_parallel_size=2,
94 data_parallel_rank=1,
95 prefill_context_parallel_size=2,
96 enable_expert_parallel=enable_ep,
97 )
98 metadata = DPMetadata.make(config, 2, torch.tensor([1, 2]))
99 manager = AgRsAll2AllManager.__new__(AgRsAll2AllManager)
100 manager.dp_world_size = 2
101 manager.use_ep = enable_ep
102
103 manager.tp_group = SimpleNamespace(world_size=1)
104 for name in ("dp", "ep"):
105 monkeypatch.setattr(
106 f"vllm.distributed.device_communicators.all2all.get_{name}_group",
107 FakeGroup,
108 )
109 monkeypatch.setattr(
110 "vllm.distributed.device_communicators.all2all.get_pcp_group",
111 lambda: SimpleNamespace(world_size=2),
112 )
113 monkeypatch.setattr(
114 "vllm.distributed.device_communicators.all2all.get_forward_context",
115 lambda: SimpleNamespace(dp_metadata=metadata),
116 )
117 with metadata.sp_local_sizes(1, pcp_size=2, use_ep=enable_ep):
118 hidden_states, _, _ = manager.dispatch(
119 torch.tensor(local_tokens),
120 torch.ones(len(local_tokens)),
121 torch.zeros(len(local_tokens)),
122 )
123 combined = manager.combine(hidden_states)
124 assert combined.tolist() == local_tokens
125 sizes = [1, 1, 2, 2] if enable_ep else [2, 4]
126 assert calls == [("gather", sizes), ("scatter", sizes)]
127
128
129 @pytest.mark.parametrize("slots", [[3, 4], [3, 4, -1, -1, -1, -1]])
130 def test_decode_cache_write_ignores_dp_padding(slots):
131 kv = torch.arange(12).reshape(6, 2)
132 pe = torch.arange(6).reshape(6, 1, 1)
133 slots = torch.tensor(slots)
134 cache_kv, cache_pe, cache_slots = maybe_gather_mla_latent_cache_inputs(
135 kv, pe, slots, num_decode_tokens=2, use_pcp=True
136 )
137 torch.testing.assert_close(cache_kv, kv[:2])
138 torch.testing.assert_close(cache_pe, pe[:2])
139 torch.testing.assert_close(cache_slots, slots[:2])
140
141
142 def test_expanded_slot_mapping_keeps_pcp_prefill_padding(monkeypatch):
143 calls = []
144
145 def all_gather(tensor, dim):
146 calls.append(tensor.shape[0])
147 return torch.cat((tensor, tensor), dim=dim)
148
149 monkeypatch.setattr(
150 "vllm.v1.attention.ops.pcp.get_pcp_group",
151 lambda: SimpleNamespace(world_size=2, all_gather=all_gather),
152 )
153 kv = torch.zeros(3, 2) # Two decodes followed by one PCP padding row.
154 pe = torch.zeros(3, 1, 1)
155 slots = torch.tensor([3, 4, 8, 3, 4, -1])
156 cache_kv, _, cache_slots = maybe_gather_mla_latent_cache_inputs(
157 kv, pe, slots, num_decode_tokens=2, use_pcp=True
158 )
159 assert calls == [1, 1]
160 assert cache_kv.shape == (4, 2)
161 assert cache_slots.tolist() == [3, 4, 8, -1]