@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 259 lines 10.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """The decoder-side SWA bounded replay batch DeepseekV41ModelState prepares."""
4
5 from dataclasses import replace
6 from types import SimpleNamespace
7 from unittest.mock import MagicMock
8
9 import numpy as np
10 import pytest
11 import torch
12
13 from vllm.config import CUDAGraphMode
14 from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
15 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
16 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
17
18 pytestmark = pytest.mark.skipif(
19 not torch.cuda.is_available(), reason="the replay buffers live on the GPU"
20 )
21
22 WINDOW = 128
23 DEVICE = torch.device("cuda")
24 GROUPS = SimpleNamespace(
25 kv_cache_groups=[
26 SimpleNamespace(
27 layer_names=["swa"],
28 kv_cache_spec=SimpleNamespace(
29 prefix_cacheable=False, prefix_replay_tokens=WINDOW
30 ),
31 ),
32 SimpleNamespace(
33 layer_names=["mla"],
34 kv_cache_spec=SimpleNamespace(
35 prefix_cacheable=True, prefix_replay_tokens=0
36 ),
37 ),
38 ]
39 )
40 BLOCK_TABLES = (torch.zeros(4, 4, device=DEVICE),) * 2
41
42 # decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
43 QUERY_LENS = [1, 300, 100]
44 SEQ_LENS = [500, 300, 300]
45 PREFILLING = [False, True, True]
46 REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
47
48
49 @pytest.fixture
50 def state(monkeypatch):
51 """A DeepseekV41ModelState whose attention builds are recorded, not run."""
52 cfg = MagicMock()
53 cfg.model_config.enable_prompt_embeds = False
54 cfg.model_config.uses_mrope = False
55 cfg.model_config.is_multimodal_model = False
56 cfg.scheduler_config.max_num_seqs = 16
57 cfg.scheduler_config.max_num_batched_tokens = 1024
58 cfg.parallel_config.data_parallel_size = 1
59 cfg.compilation_config.fast_moe_cold_start = False
60 layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None)
61 model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers)
62 builds: list = []
63
64 def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
65 build = SimpleNamespace(
66 batch=input_batch,
67 cg_mode=cg_mode,
68 slot_mappings=slot_mappings.clone(),
69 replay_start=kw["model_specific_attn_metadata"].replay_start,
70 attn_metadata=object(),
71 )
72 builds.append(build)
73 return {"swa": build.attn_metadata}
74
75 monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
76 state = DeepseekV41ModelState(cfg, model, None, DEVICE)
77 state.builds = builds # type: ignore[attr-defined]
78 return state
79
80
81 def _input_batch(
82 query_lens: list[int],
83 seq_lens: list[int],
84 is_prefilling: list[bool],
85 device_query_lens: list[int] | None = None,
86 num_tokens_after_padding: int | None = None,
87 max_query_len: int | None = None,
88 ) -> InputBatch:
89 num_reqs, num_tokens = len(query_lens), sum(query_lens)
90 num_padded = num_tokens_after_padding or num_tokens
91 batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(16, 1024, DEVICE))
92 return replace(
93 batch,
94 num_tokens=num_tokens,
95 num_tokens_after_padding=num_padded,
96 query_start_loc=torch.tensor(
97 [0, *np.cumsum(device_query_lens or query_lens)],
98 dtype=torch.int32,
99 device=DEVICE,
100 ),
101 query_start_loc_np=np.array([0, *np.cumsum(query_lens)], dtype=np.int32),
102 seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE),
103 seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
104 is_prefilling_np=np.array(is_prefilling, dtype=np.bool_),
105 positions=torch.cat(
106 [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)]
107 + [torch.zeros(num_padded - num_tokens, dtype=torch.int64)]
108 ).to(DEVICE),
109 is_padding=torch.arange(num_padded, device=DEVICE) >= num_tokens,
110 max_query_len=max_query_len,
111 )
112
113
114 def _slot_mappings(num_tokens: int) -> torch.Tensor:
115 return torch.stack([torch.arange(num_tokens), torch.arange(num_tokens) * 10]).to(
116 DEVICE
117 )
118
119
120 def _prepare(state, batch, cg_mode):
121 """Runs prepare_attn; returns the replay batch and the replay build's inputs."""
122 state.prepare_attn(
123 batch,
124 cg_mode,
125 BLOCK_TABLES,
126 _slot_mappings(batch.num_tokens_after_padding),
127 [],
128 GROUPS,
129 )
130 layers = state.decoder_replay_layers
131 if layers.rows is None:
132 assert layers.forward_context is None
133 return None, None
134 return SimpleNamespace(
135 rows=layers.rows, forward_context=layers.forward_context
136 ), state.builds[-1]
137
138
139 def test_replay_batch_keeps_each_request_window(state):
140 state._replay_start_np[2] = 50 # the encoder-side replay start of request 2
141 batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
142 replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
143 assert replay is not None
144 rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
145 assert torch.equal(replay.rows, rows)
146 sub = build.batch
147 assert build.cg_mode == CUDAGraphMode.NONE
148 assert sub.query_start_loc.tolist() == [0, 1, 129, 229]
149 assert sub.query_start_loc_np.tolist() == [0, 1, 129, 229]
150 assert sub.num_tokens == sub.num_tokens_after_padding == 229
151 assert sub.max_query_len == WINDOW
152 # The trimmed request's window starts at its replay window; a higher
153 # encoder-side replay start stands.
154 assert build.replay_start.tolist() == [0, 300 - WINDOW, 50]
155 assert torch.equal(sub.positions, batch.positions[rows])
156 assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows])
157 context = replay.forward_context
158 assert context.attn_metadata == {"swa": build.attn_metadata}
159 assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows])
160 assert context.dp_metadata is None and context.is_padding is None
161
162
163 def test_replay_batch_keeps_device_decode_boundaries(state):
164 """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device
165 [1, 3]): their rows stay and the boundaries shift only by trimmed rows."""
166 batch = _input_batch(
167 [2, 2, 300], [10, 10, 300], PREFILLING, device_query_lens=[1, 3, 300]
168 )
169 replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
170 assert replay is not None
171 assert replay.rows.tolist() == [0, 1, 2, 3, *range(304 - WINDOW, 304)]
172 assert build.batch.query_start_loc.tolist() == [0, 1, 4, 4 + WINDOW]
173
174
175 def test_prompt_logprobs_requests_keep_their_rows(state):
176 """Every prompt row of a prompt-logprobs request is read, so it never
177 trims; the other prefills still do."""
178 state._keeps_rows_np[1] = True
179 batch = _input_batch([1, 300, 300], [500, 300, 300], PREFILLING)
180 replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
181 assert replay is not None
182 assert replay.rows.tolist() == [0, *range(1, 301), *range(601 - WINDOW, 601)]
183 assert build.batch.max_query_len == 300
184
185
186 def test_replay_batch_keeps_adaptive_verification_query_bound(state):
187 """Adaptive verification bounds the decodes' device-side query lengths
188 above their CPU lengths; the replay batch keeps that bound."""
189 batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, max_query_len=200)
190 replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
191 assert replay is not None and build.batch.max_query_len == 200
192
193
194 def test_only_eager_steps_trim(state):
195 """A CUDA graph keeps the layers on its whole batch; an eager step trims a
196 padded batch to its real rows."""
197 batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512)
198 for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
199 assert _prepare(state, batch, cg_mode) == (None, None)
200 replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
201 assert replay is not None
202 assert replay.rows.tolist() == REPLAY_ROWS
203 assert build.cg_mode == CUDAGraphMode.NONE
204 assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS)
205
206
207 def test_whole_batch_forwards_get_no_replay_batch(state):
208 short = _input_batch([100, 100], [100, 100], [True, True])
209 for cg_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
210 assert _prepare(state, short, cg_mode) == (None, None)
211 assert len(state.builds) == 3 # the batch's own metadata only
212 # Dummy batches, captures included (prepared as eager), are not prefills and
213 # keep their rows.
214 dummy = _input_batch(
215 QUERY_LENS, SEQ_LENS, [False] * 3, num_tokens_after_padding=512
216 )
217 assert _prepare(state, dummy, CUDAGraphMode.NONE) == (None, None)
218
219
220 @pytest.fixture
221 def dp_state(state, monkeypatch):
222 """`state` on DP rank 0 of 2; `state.other` sets what rank 1 reports
223 (trims, replay tokens)."""
224 import vllm.models.deepseek_v41.nvidia.model_state as module
225
226 state.vllm_config.parallel_config.data_parallel_size = 2
227 state.vllm_config.parallel_config.data_parallel_rank = 0
228 state.vllm_config.parallel_config.is_moe_model = True
229 state.other = (False, 0) # type: ignore[attr-defined]
230
231 def all_reduce(tensor, group):
232 tensor[:, 1] = torch.tensor(state.other, dtype=torch.int32)
233
234 monkeypatch.setattr(module.dist, "all_reduce", all_reduce)
235 monkeypatch.setattr(module, "get_dp_group", lambda: SimpleNamespace(cpu_group=None))
236 return state
237
238
239 def test_dp_ranks_replay_together(dp_state):
240 """A trimming rank makes the others replay too, on every rank's count."""
241 short = _input_batch([100, 100], [100, 100], [True, True])
242 dp_state.other = (False, 200)
243 assert _prepare(dp_state, short, CUDAGraphMode.NONE) == (None, None)
244
245 dp_state.other = (True, 300) # rank 1 trims to 300 rows
246 replay, build = _prepare(dp_state, short, CUDAGraphMode.NONE)
247 assert replay is not None
248 assert replay.rows.tolist() == list(range(200)) # this rank keeps every row
249 dp_metadata = replay.forward_context.dp_metadata
250 assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [200, 300]
251
252
253 def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
254 dummy = _input_batch([512], [512], [False])
255 dp_state.other = (True, 129)
256 replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE)
257 assert replay is not None and replay.rows.shape[0] == WINDOW
258 dp_metadata = replay.forward_context.dp_metadata
259 assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129]