@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 214 lines 7.52 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5 from unittest.mock import patch
6
7 import numpy as np
8 import pytest
9 import torch
10
11 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
13 DFlashSpeculator,
14 )
15
16 pytestmark = pytest.mark.skipif(
17 not torch.cuda.is_available(), reason="requires a CUDA device"
18 )
19
20
21 def _run_prepare(
22 *,
23 target_positions: list[int],
24 block_table_values: list[int],
25 cp_rank: int = 0,
26 cp_size: int = 1,
27 cp_interleave: int = 1,
28 draft_dcp_size: int | None = None,
29 ):
30 device = torch.device("cuda")
31 max_num_reqs = 4
32 max_num_tokens = 16
33 num_speculative_steps = 3
34
35 input_buffers = SimpleNamespace(
36 input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
37 positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
38 query_start_loc=torch.full(
39 (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
40 ),
41 seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device),
42 )
43 input_batch = SimpleNamespace(
44 num_reqs=1,
45 num_tokens=4,
46 num_scheduled_tokens=np.array([4], dtype=np.int32),
47 seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]),
48 positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
49 query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
50 idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
51 )
52 query_slot_mapping = torch.full(
53 (max_num_tokens,), -2, dtype=torch.int64, device=device
54 )
55 context_positions = torch.full(
56 (max_num_tokens,), -1, dtype=torch.int64, device=device
57 )
58 context_slot_mapping = torch.full(
59 (max_num_tokens,), -2, dtype=torch.int64, device=device
60 )
61 sample_indices = torch.full(
62 (max_num_reqs * num_speculative_steps,),
63 -1,
64 dtype=torch.int64,
65 device=device,
66 )
67 sample_pos = torch.full_like(sample_indices, -1)
68 sample_idx_mapping = torch.full(
69 sample_indices.shape, -1, dtype=torch.int32, device=device
70 )
71 temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
72 seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
73 input_temperature = torch.tensor(
74 [0.0, 0.0, 1.0, 0.0], dtype=torch.float32, device=device
75 )
76 input_seeds = torch.tensor([0, 0, 17, 0], dtype=torch.int64, device=device)
77 last_sampled = torch.tensor([0, 0, 99, 0], dtype=torch.int64, device=device)
78 next_prefill_tokens = torch.zeros_like(last_sampled)
79 block_table = torch.tensor([block_table_values], dtype=torch.int32, device=device)
80
81 class InputsPrepared(Exception):
82 pass
83
84 draft = SimpleNamespace(
85 input_buffers=input_buffers,
86 context_positions=context_positions,
87 _context_slot_mappings=context_slot_mapping.unsqueeze(0),
88 sample_indices=sample_indices,
89 sample_pos=sample_pos,
90 sample_idx_mapping=sample_idx_mapping,
91 temperature=temperature,
92 seeds=seeds,
93 hidden_states=torch.zeros(4, 1, device=device),
94 prepare_context_anchor=lambda *args: None,
95 query_cudagraph_manager=None,
96 dp_size=1,
97 dp_rank=0,
98 pcp_manager=None,
99 draft_kv_cache_group_id=0,
100 draft_kv_cache_group_ids=[0],
101 _layer_group_idx=None,
102 block_tables=SimpleNamespace(
103 slot_mappings=query_slot_mapping.unsqueeze(0),
104 input_block_tables=[block_table],
105 kernel_block_sizes=[4],
106 cp_rank=cp_rank,
107 cp_size=cp_size,
108 cp_interleave=cp_interleave,
109 ),
110 dcp_size=cp_size if draft_dcp_size is None else draft_dcp_size,
111 parallel_drafting_token_id=123,
112 num_query_per_req=num_speculative_steps,
113 num_speculative_steps=num_speculative_steps,
114 max_num_reqs=max_num_reqs,
115 max_num_tokens=max_num_tokens,
116 max_model_len=128,
117 sample_from_anchor=True,
118 )
119 with (
120 patch(
121 "vllm.v1.worker.gpu.spec_decode.dflash.speculator.dispatch_cg_and_sync_dp",
122 side_effect=InputsPrepared,
123 ),
124 pytest.raises(InputsPrepared),
125 ):
126 DFlashSpeculator.propose(
127 draft,
128 input_batch,
129 {},
130 {},
131 draft.hidden_states,
132 None,
133 torch.tensor([1], dtype=torch.int32, device=device),
134 torch.tensor([2], dtype=torch.int32, device=device),
135 last_sampled,
136 next_prefill_tokens,
137 input_temperature,
138 input_seeds,
139 )
140 torch.accelerator.synchronize()
141 return SimpleNamespace(
142 input_buffers=input_buffers,
143 query_slot_mapping=query_slot_mapping.cpu(),
144 context_positions=context_positions.cpu(),
145 context_slot_mapping=context_slot_mapping.cpu(),
146 sample_indices=sample_indices.cpu(),
147 sample_pos=sample_pos.cpu(),
148 sample_idx_mapping=sample_idx_mapping.cpu(),
149 temperature=temperature.cpu(),
150 seeds=seeds.cpu(),
151 )
152
153
154 @pytest.mark.parametrize("cp_rank,cp_size", [(0, 1), (0, 2), (1, 2)])
155 def test_prepare_dflash_inputs_excludes_rejected_context_suffix(cp_rank, cp_size):
156 # Positions 10/11 use physical block 7. Rejected positions 12/13 would use
157 # block 8, but must be PAD context rather than contaminating draft KV.
158 out = _run_prepare(
159 target_positions=[10, 11, 12, 13],
160 block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
161 cp_rank=cp_rank,
162 cp_size=cp_size,
163 cp_interleave=2,
164 draft_dcp_size=1,
165 )
166
167 assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
168 assert out.context_slot_mapping[:4].tolist() == [30, 31, PAD_SLOT_ID, PAD_SLOT_ID]
169
170 # The replacement query starts immediately after the two valid rows and
171 # advances from the last accepted position (11).
172 assert out.input_buffers.input_ids[:3].cpu().tolist() == [99, 123, 123]
173 assert out.input_buffers.positions[:3].cpu().tolist() == [12, 13, 14]
174 assert out.query_slot_mapping[:3].tolist() == [32, 33, 34]
175 assert out.sample_indices[:3].tolist() == [0, 1, 2]
176 assert out.sample_pos[:3].tolist() == [13, 14, 15]
177 assert out.sample_idx_mapping[:3].tolist() == [2, 2, 2]
178 assert out.temperature[2].item() == 1.0
179 assert out.seeds[2].item() == 17
180
181
182 def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
183 out = _run_prepare(
184 target_positions=[10, 11, 12, 13],
185 block_table_values=[0, 7, 8, 9],
186 cp_rank=1,
187 cp_size=2,
188 cp_interleave=2,
189 )
190
191 assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
192 assert out.context_slot_mapping[:4].tolist() == [28, 29, PAD_SLOT_ID, PAD_SLOT_ID]
193 assert out.query_slot_mapping[:3].tolist() == [PAD_SLOT_ID, PAD_SLOT_ID, 30]
194
195
196 def test_prepare_dflash_inputs_never_writes_the_null_block():
197 # The valid context uses logical block 0 and the replacement query uses
198 # logical block 1. Both map to the null block and must remain unwritable.
199 out = _run_prepare(
200 target_positions=[2, 3, 4, 5],
201 block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
202 )
203
204 assert out.context_slot_mapping[:4].tolist() == [
205 PAD_SLOT_ID,
206 PAD_SLOT_ID,
207 PAD_SLOT_ID,
208 PAD_SLOT_ID,
209 ]
210 assert out.query_slot_mapping[:3].tolist() == [
211 PAD_SLOT_ID,
212 PAD_SLOT_ID,
213 PAD_SLOT_ID,
214 ]