@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 189 lines 7.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import cast
5 from unittest.mock import Mock
6
7 import pytest
8 import torch
9
10 import vllm.v1.worker.workspace as workspace
11 from vllm.config import VllmConfig
12 from vllm.v1.worker.gpu_worker import _num_workspace_lanes
13
14
15 class _SpecConfig:
16 def __init__(self, dspark: bool) -> None:
17 self._dspark = dspark
18
19 def use_dspark(self) -> bool:
20 return self._dspark
21
22
23 class _VllmConfig:
24 def __init__(self, spec_config: _SpecConfig | None) -> None:
25 self.speculative_config = spec_config
26
27
28 @pytest.mark.parametrize(
29 ("use_v2_model_runner", "spec_config", "expected"),
30 [
31 (True, _SpecConfig(True), 2),
32 (False, _SpecConfig(True), 1),
33 (True, _SpecConfig(False), 1),
34 (True, None, 1),
35 ],
36 )
37 def test_workspace_lane_count_is_dspark_only(
38 use_v2_model_runner: bool,
39 spec_config: _SpecConfig | None,
40 expected: int,
41 ) -> None:
42 config = cast(VllmConfig, _VllmConfig(spec_config))
43 assert _num_workspace_lanes(config, use_v2_model_runner) == expected
44
45
46 def test_workspace_lanes_do_not_alias_and_restore_context(monkeypatch) -> None:
47 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
48 manager = workspace.WorkspaceManager(
49 torch.device("cpu"), num_ubatches=2, num_lanes=2
50 )
51
52 assert manager._current_workspaces == [None, None, None, None]
53
54 (target,) = manager.get_simultaneous(((512,), torch.uint8))
55 with workspace.use_workspace_lane(1):
56 (draft,) = manager.get_simultaneous(((256,), torch.uint8))
57 (draft_reused,) = manager.get_simultaneous(((8,), torch.uint8))
58 (target_reused,) = manager.get_simultaneous(((8,), torch.uint8))
59
60 assert manager._current_workspaces[0].numel() == 512 # type: ignore[union-attr]
61 assert manager._current_workspaces[1].numel() == 256 # type: ignore[union-attr]
62 assert manager._current_workspaces[2:] == [None, None]
63 assert target.data_ptr() != draft.data_ptr()
64 assert draft.data_ptr() == draft_reused.data_ptr()
65 assert target.data_ptr() == target_reused.data_ptr()
66
67
68 def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None:
69 active_ubatch = [0]
70 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0])
71 manager = workspace.WorkspaceManager(
72 torch.device("cpu"), num_ubatches=2, num_lanes=2
73 )
74
75 pointers = set()
76 for ubatch_id in range(2):
77 active_ubatch[0] = ubatch_id
78 for lane in range(2):
79 with workspace.use_workspace_lane(lane):
80 (buffer,) = manager.get_simultaneous(((16,), torch.uint8))
81 pointers.add(buffer.data_ptr())
82
83 assert len(pointers) == 4
84
85
86 def test_workspace_lock_blocks_growth_and_unlock_restores(monkeypatch) -> None:
87 """Once locked, oversized requests fail loudly instead of reallocating the
88 buffer that captured CUDA graphs point at; unlock restores growth."""
89 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
90 manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1)
91
92 (buf,) = manager.get_simultaneous(((256,), torch.uint8))
93 manager.lock()
94 assert manager.is_locked()
95
96 # Requests within the reserved size still reuse the same buffer.
97 (same,) = manager.get_simultaneous(((256,), torch.uint8))
98 (smaller,) = manager.get_simultaneous(((8,), torch.uint8))
99 assert same.data_ptr() == buf.data_ptr()
100 assert smaller.data_ptr() == buf.data_ptr()
101
102 with pytest.raises(AssertionError, match="Workspace is locked"):
103 manager.get_simultaneous(((512,), torch.uint8))
104
105 manager.unlock()
106 (grown,) = manager.get_simultaneous(((512,), torch.uint8))
107 assert grown.numel() == 512
108
109
110 def test_workspace_lane_validation(monkeypatch) -> None:
111 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
112 manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1)
113
114 with (
115 pytest.raises(ValueError, match="non-negative"),
116 workspace.use_workspace_lane(-1),
117 ):
118 pass
119
120 with (
121 workspace.use_workspace_lane(1),
122 pytest.raises(RuntimeError, match="is not configured"),
123 ):
124 manager.get_simultaneous(((1,), torch.uint8))
125
126 with pytest.raises(ValueError, match="at least one"):
127 workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0)
128
129
130 def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None:
131 """Only the requesting slot allocates; each ubatch/lane owns its resource."""
132 active_ubatch = 0
133 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch)
134 manager = workspace.WorkspaceManager(
135 torch.device("cpu"), num_ubatches=2, num_lanes=2
136 )
137 factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)})
138 resources = []
139 for ubatch in range(2):
140 active_ubatch = ubatch
141 for lane in range(2):
142 with workspace.use_workspace_lane(lane):
143 resource = manager.get_persistent_resource("scratch", factory)
144 resources.append(resource)
145 resource["buffer"].fill_(len(resources))
146 assert manager.get_persistent_resource("scratch", factory) is resource
147 assert factory.call_count == len(resources)
148
149 assert len({r["buffer"].data_ptr() for r in resources}) == 4
150 for i, resource in enumerate(resources, 1):
151 assert torch.all(resource["buffer"] == i)
152
153
154 def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None:
155 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
156 manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2)
157 factory = Mock(side_effect=object)
158 first = manager.get_persistent_resource("scratch", factory)
159 manager.lock()
160 assert manager.get_persistent_resource("scratch", factory) is first
161 with pytest.raises(AssertionError, match="was not allocated during warmup"):
162 manager.get_persistent_resource("new", factory)
163 with workspace.use_workspace_lane(1):
164 with pytest.raises(AssertionError, match="was not allocated during warmup"):
165 manager.get_persistent_resource("scratch", factory)
166 assert factory.call_count == 1
167 manager.unlock()
168 assert manager.get_persistent_resource("scratch", factory) is not first
169 assert manager.get_persistent_resource("scratch", factory) is first
170
171
172 def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None:
173 """Initialization happens once and transient scratch cannot overwrite locks."""
174 manager = workspace.WorkspaceManager(torch.device("cpu"))
175 first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True)
176 assert torch.count_nonzero(first) == 0
177 first.fill_(1)
178 (scratch,) = manager.get_simultaneous(((1024,), torch.uint8))
179 scratch.zero_()
180 assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first
181 assert torch.all(first == 1)
182 with pytest.raises(ValueError, match="requested shape"):
183 manager.get_persistent("locks", (16,), torch.int32)
184 with pytest.raises(ValueError, match="requested shape"):
185 manager.get_persistent("locks", (8,), torch.int64)
186 manager.lock()
187 assert manager.get_persistent("locks", (8,), torch.int32) is first
188 with pytest.raises(AssertionError, match="was not allocated during warmup"):
189 manager.get_persistent("new", (8,), torch.int32)