| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import nullcontext |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.config.compilation import CUDAGraphMode |
| 11 | from vllm.v1.attention.backend import AttentionCGSupport, AttentionMetadataBuilder |
| 12 | from vllm.v1.worker.gpu.async_utils import StepTimingSample |
| 13 | from vllm.v1.worker.gpu.spec_decode import adaptive_verification as adaptive_module |
| 14 | from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( |
| 15 | AdaptiveVerificationManager, |
| 16 | maybe_create_adaptive_verification_manager, |
| 17 | resolve_adaptive_cudagraph_mode, |
| 18 | ) |
| 19 | from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping |
| 20 | |
| 21 | |
| 22 | def make_manager( |
| 23 | confidences: np.ndarray, verify_cost_ms: np.ndarray |
| 24 | ) -> AdaptiveVerificationManager: |
| 25 | num_reqs, num_steps = confidences.shape |
| 26 | manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager) |
| 27 | manager.num_speculative_steps = num_steps |
| 28 | manager._stale_confidences = [SimpleNamespace(np=confidences)] # type: ignore[list-item] # CPU-only confidence buffer. |
| 29 | manager._stale_idx = 0 |
| 30 | manager.req_states = SimpleNamespace( |
| 31 | req_id_to_index={"low": 0, "high": 1}, |
| 32 | num_computed_tokens_np=np.ones(num_reqs, dtype=np.int32), |
| 33 | prefill_len=SimpleNamespace(np=np.ones(num_reqs, dtype=np.int32)), |
| 34 | ) |
| 35 | manager.cost_tables = (np.zeros(num_reqs + 1), verify_cost_ms) |
| 36 | manager._max_total_logits = 1 << 30 |
| 37 | manager.num_bonus_tokens = 1 |
| 38 | return manager |
| 39 | |
| 40 | |
| 41 | @pytest.mark.parametrize( |
| 42 | ("mode", "piecewise_capture_available", "expected"), |
| 43 | [ |
| 44 | ("FULL_DECODE_ONLY", True, "FULL_DECODE_ONLY"), |
| 45 | ("FULL_DECODE_ONLY", False, "FULL_DECODE_ONLY"), |
| 46 | ("FULL", True, "FULL_AND_PIECEWISE"), |
| 47 | ("FULL", False, "FULL_DECODE_ONLY"), |
| 48 | ("FULL_AND_PIECEWISE", True, "FULL_AND_PIECEWISE"), |
| 49 | ("FULL_AND_PIECEWISE", False, "FULL_DECODE_ONLY"), |
| 50 | ], |
| 51 | ) |
| 52 | def test_resolve_adaptive_cudagraph_mode(mode, piecewise_capture_available, expected): |
| 53 | assert ( |
| 54 | resolve_adaptive_cudagraph_mode( |
| 55 | CUDAGraphMode[mode], piecewise_capture_available=piecewise_capture_available |
| 56 | ) |
| 57 | == CUDAGraphMode[expected] |
| 58 | ) |
| 59 | |
| 60 | |
| 61 | @pytest.mark.parametrize( |
| 62 | "target_support,target_bound,device_offsets,additional,error", |
| 63 | [ |
| 64 | pytest.param(AttentionCGSupport.ALWAYS, None, True, None, None, id="always"), |
| 65 | pytest.param( |
| 66 | AttentionCGSupport.UNIFORM_BATCH, 8, True, None, None, id="bounded" |
| 67 | ), |
| 68 | pytest.param( |
| 69 | AttentionCGSupport.UNIFORM_BATCH, |
| 70 | None, |
| 71 | True, |
| 72 | None, |
| 73 | "TargetBackend allows none", |
| 74 | id="unsupported", |
| 75 | ), |
| 76 | pytest.param( |
| 77 | AttentionCGSupport.UNIFORM_BATCH, |
| 78 | 7, |
| 79 | True, |
| 80 | None, |
| 81 | "up to 8, but TargetBackend allows at most 7", |
| 82 | id="too-narrow", |
| 83 | ), |
| 84 | pytest.param( |
| 85 | AttentionCGSupport.UNIFORM_BATCH, |
| 86 | 8, |
| 87 | False, |
| 88 | None, |
| 89 | "trims verification requests", |
| 90 | id="host-query-lens", |
| 91 | ), |
| 92 | pytest.param( |
| 93 | AttentionCGSupport.ALWAYS, |
| 94 | None, |
| 95 | True, |
| 96 | (AttentionCGSupport.UNIFORM_BATCH, "EncoderBackend"), |
| 97 | "EncoderBackend allows none", |
| 98 | id="additional-group", |
| 99 | ), |
| 100 | ], |
| 101 | ) |
| 102 | def test_manager_checks_target_varlen_cudagraph_bound( |
| 103 | monkeypatch, target_support, target_bound, device_offsets, additional, error |
| 104 | ): |
| 105 | """Target builders must replay varlen decode graphs of the verification |
| 106 | width; draft-only groups are not checked.""" |
| 107 | |
| 108 | def group(backend_name, layer_name, support, bound=None, mismatch=True): |
| 109 | class Builder(AttentionMetadataBuilder): |
| 110 | _cudagraph_support = support |
| 111 | |
| 112 | def __init__(self): |
| 113 | pass |
| 114 | |
| 115 | def build(self, common_prefix_len, common_attn_metadata, fast_build=False): |
| 116 | raise NotImplementedError |
| 117 | |
| 118 | if bound is not None: |
| 119 | |
| 120 | @classmethod |
| 121 | def get_varlen_cudagraph_max_query_len(cls, *_args): |
| 122 | return bound |
| 123 | |
| 124 | backend = type( |
| 125 | backend_name, |
| 126 | (), |
| 127 | {"supports_device_cpu_query_lens_mismatch": staticmethod(lambda: mismatch)}, |
| 128 | ) |
| 129 | builder = Builder() |
| 130 | return SimpleNamespace( |
| 131 | layer_names=[layer_name], |
| 132 | backend=backend, |
| 133 | kv_cache_spec=None, |
| 134 | get_metadata_builder=lambda _index: builder, |
| 135 | ) |
| 136 | |
| 137 | groups = [ |
| 138 | [ |
| 139 | group( |
| 140 | "TargetBackend", "target", target_support, target_bound, device_offsets |
| 141 | ), |
| 142 | group("DraftBackend", "draft", AttentionCGSupport.UNIFORM_BATCH), |
| 143 | ] |
| 144 | ] |
| 145 | created = object() |
| 146 | monkeypatch.setattr( |
| 147 | adaptive_module, |
| 148 | "AdaptiveVerificationManager", |
| 149 | lambda *_args, **_kwargs: created, |
| 150 | ) |
| 151 | |
| 152 | with pytest.raises(ValueError, match=error) if error else nullcontext(): |
| 153 | manager = maybe_create_adaptive_verification_manager( |
| 154 | enable_adaptive_verification=True, |
| 155 | attn_groups=groups, |
| 156 | req_states=SimpleNamespace(num_speculative_steps=7), |
| 157 | query_start_loc=object(), |
| 158 | num_bonus_tokens=1, |
| 159 | max_total_logits=1, |
| 160 | vllm_config=None, |
| 161 | target_layer_names={"target"}, |
| 162 | additional_attn_cg_support=additional, |
| 163 | ) |
| 164 | assert manager is created |
| 165 | |
| 166 | |
| 167 | def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves(): |
| 168 | # Verification is cheap up to two extra tokens, then jumps 100x; only the |
| 169 | # highest-confidence draft is worth the cheap slot. |
| 170 | manager = make_manager( |
| 171 | np.array([[0.1, 0.1], [0.9, 0.9]], dtype=np.float32), |
| 172 | np.array([1.0, 1.0, 1.0, 1.0, 100.0, 100.0, 100.0]), |
| 173 | ) |
| 174 | |
| 175 | manager.get_num_tokens( |
| 176 | {"low": 3, "high": 3}, |
| 177 | {"low": [1, 2], "high": [3, 4]}, |
| 178 | ) |
| 179 | assert manager._batch_budget is not None |
| 180 | valid_drafts, num_non_draft_tokens, draft_budget = manager._batch_budget |
| 181 | |
| 182 | assert draft_budget == 1 |
| 183 | assert valid_drafts == {"low": 2, "high": 2} |
| 184 | assert num_non_draft_tokens == {"low": 1, "high": 1} |
| 185 | |
| 186 | |
| 187 | def test_profiled_batches_seed_cost_curves_via_consumer(monkeypatch): |
| 188 | manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager) |
| 189 | manager.req_states = SimpleNamespace(max_num_batched_tokens=4096, max_num_reqs=64) |
| 190 | manager.num_speculative_steps = 7 |
| 191 | manager.num_bonus_tokens = 1 |
| 192 | curves: dict[str, list[tuple[int, float]]] = {} |
| 193 | monkeypatch.setattr( |
| 194 | manager, |
| 195 | "set_cost_curves", |
| 196 | lambda draft, verify: curves.update(draft=draft, verify=verify), |
| 197 | ) |
| 198 | |
| 199 | timings = [ |
| 200 | StepTimingSample( |
| 201 | forward_ms=float(batch["num_tokens"]), |
| 202 | drafter_ms=1.0, |
| 203 | num_target_tokens=batch["num_tokens"], |
| 204 | num_reqs=batch["num_tokens"] // 8, |
| 205 | # Only the captured sizes replay a graph; the tail sizes run eager. |
| 206 | full_cudagraph=batch["num_tokens"] <= 1024, |
| 207 | ) |
| 208 | for batch in manager.batches_to_profile([8, 1024]) |
| 209 | ] |
| 210 | manager.set_initial_cost_curves(timings) |
| 211 | |
| 212 | # Tail beyond the last capture size: 1.5x then doubling to the max. |
| 213 | assert curves["verify"] == [ |
| 214 | (8, 8.0), |
| 215 | (1024, 1024.0), |
| 216 | (1536, 1536.0), |
| 217 | (2048, 2048.0), |
| 218 | (4096, 4096.0), |
| 219 | ] |
| 220 | # Eager batches must not contribute to the draft curve: keyed by request |
| 221 | # count they would land inside the captured range and, once made monotonic, |
| 222 | # smear that eager cost across every larger request count. |
| 223 | assert curves["draft"] == [(1, 1.0), (128, 1.0)] |
| 224 | |
| 225 | |
| 226 | def test_compact_batch_preserves_totals_and_bounds(): |
| 227 | # The CPU placeholder layout must keep the batch total equal to the GPU |
| 228 | # total and every verification row within decode_query_len, or downstream |
| 229 | # CPU metadata desyncs from the reallocated GPU boundaries. |
| 230 | manager = make_manager( |
| 231 | np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), |
| 232 | np.array([1.0] * 44 + [100.0] * 3), |
| 233 | ) |
| 234 | manager.req_states.req_id_to_index["prefill"] = 2 |
| 235 | manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) |
| 236 | manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) |
| 237 | num_tokens = manager.get_num_tokens( |
| 238 | {"low": 3, "high": 3, "prefill": 40}, |
| 239 | {"low": [1, 2], "high": [3, 4]}, |
| 240 | ) |
| 241 | scheduled = np.array([3, 3, 40], dtype=np.int32) |
| 242 | drafts = np.array([2, 2, 0], dtype=np.int32) |
| 243 | cu_num_logits_np = np.array([0, 3, 6, 7], dtype=np.int32) |
| 244 | compacted, _ = manager.compact_batch(drafts, scheduled, cu_num_logits_np) |
| 245 | |
| 246 | assert int(compacted.sum()) == num_tokens |
| 247 | num_steps = manager.num_speculative_steps |
| 248 | assert (compacted[:2] <= 1 + num_steps).all() |
| 249 | assert compacted[2] == 40 |
| 250 | |
| 251 | |
| 252 | def test_budget_caps_at_one_rejection_sampler_chunk(): |
| 253 | # The chunked verification path cannot address the compacted logits |
| 254 | # layout, so the budget must keep total logits within a single chunk. |
| 255 | manager = make_manager( |
| 256 | np.array([[0.9, 0.9], [0.9, 0.9]], dtype=np.float32), |
| 257 | np.ones(7), |
| 258 | ) |
| 259 | manager._max_total_logits = 3 # 2 bonus logits + at most 1 draft |
| 260 | manager.get_num_tokens( |
| 261 | {"low": 3, "high": 3}, |
| 262 | {"low": [1, 2], "high": [3, 4]}, |
| 263 | ) |
| 264 | assert manager._batch_budget is not None |
| 265 | _, _, draft_budget = manager._batch_budget |
| 266 | assert draft_budget <= 1 |
| 267 | |
| 268 | |
| 269 | def test_zero_budget_rebuilds_cpu_cu_num_logits(): |
| 270 | # When one bonus row per request already overflows a verification chunk, the |
| 271 | # budget clamps to zero but the batch still needs chunking. Every capacity is |
| 272 | # zeroed on device, so the CPU can name that layout exactly -- and must, since |
| 273 | # _iter_request_chunks slices the compacted logits with these offsets. |
| 274 | # |
| 275 | # The third request is a chunked prefill (no drafts, still mid-prompt). The |
| 276 | # runner gives *every* request num_bonus_tokens logits rows regardless |
| 277 | # (num_logits = num_draft_tokens_per_req + num_bonus_tokens), so the rebuilt |
| 278 | # offsets stay uniform rather than skipping non-verification rows. |
| 279 | manager = make_manager( |
| 280 | np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), |
| 281 | np.ones(64), |
| 282 | ) |
| 283 | manager.req_states.req_id_to_index["prefill"] = 2 |
| 284 | manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) |
| 285 | manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) |
| 286 | manager._max_total_logits = 2 # < 3 requests * 1 bonus token |
| 287 | |
| 288 | manager.get_num_tokens( |
| 289 | {"low": 3, "high": 3, "prefill": 40}, |
| 290 | {"low": [1, 2], "high": [3, 4]}, |
| 291 | ) |
| 292 | assert manager._batch_budget is not None |
| 293 | _, _, draft_budget = manager._batch_budget |
| 294 | assert draft_budget == 0 |
| 295 | |
| 296 | scheduled = np.array([3, 3, 40], dtype=np.int32) |
| 297 | drafts = np.array([2, 2, 0], dtype=np.int32) |
| 298 | scheduled_cu_num_logits = np.array([0, 3, 6, 7], dtype=np.int32) |
| 299 | compacted, cu_num_logits_np = manager.compact_batch( |
| 300 | drafts, scheduled, scheduled_cu_num_logits |
| 301 | ) |
| 302 | |
| 303 | # One bonus row per request, matching cumsum(capacities + num_bonus_tokens) |
| 304 | # with every capacity zeroed -- the prefill row included. |
| 305 | expected = np.arange(4, dtype=np.int32) * manager.num_bonus_tokens |
| 306 | assert np.array_equal(cu_num_logits_np, expected) |
| 307 | assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype |
| 308 | # The prefill keeps its scheduled tokens; only drafts are dropped. |
| 309 | assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) |
| 310 | |
| 311 | |
| 312 | def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): |
| 313 | # The scheduler sizes the grammar bitmask from the *scheduled* drafts |
| 314 | # (len(drafts) + 1 rows per request), but a zero budget rewrites |
| 315 | # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping |
| 316 | # from those rewritten offsets drops rows and trips the |
| 317 | # `num_masks == len(mapping)` assert in apply_grammar_bitmask. |
| 318 | manager = make_manager( |
| 319 | np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), |
| 320 | np.ones(64), |
| 321 | ) |
| 322 | manager.req_states.req_id_to_index["prefill"] = 2 |
| 323 | manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) |
| 324 | manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) |
| 325 | manager._max_total_logits = 2 # < 3 requests * 1 bonus token |
| 326 | |
| 327 | scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]} |
| 328 | manager.get_num_tokens( |
| 329 | {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens |
| 330 | ) |
| 331 | assert manager._batch_budget is not None |
| 332 | assert manager._batch_budget[2] == 0 |
| 333 | |
| 334 | req_ids = ["low", "high", "prefill"] |
| 335 | num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) |
| 336 | _, cu_num_logits_np = manager.compact_batch( |
| 337 | num_draft_tokens_per_req, |
| 338 | np.array([3, 3, 40], dtype=np.int32), |
| 339 | np.array([0, 3, 6, 7], dtype=np.int32), |
| 340 | ) |
| 341 | |
| 342 | mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens |
| 343 | mapping = _build_grammar_mapping( |
| 344 | req_ids, |
| 345 | req_ids, |
| 346 | cu_num_logits_np, |
| 347 | num_draft_tokens_per_req, |
| 348 | manager.num_bonus_tokens, |
| 349 | mask_stride, |
| 350 | ) |
| 351 | |
| 352 | num_bitmask_rows = sum( |
| 353 | len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids |
| 354 | ) |
| 355 | assert len(mapping) == num_bitmask_rows |
| 356 | # (request, position) keys, so the kernel can mask rows the compacted |
| 357 | # device layout no longer has room for. |
| 358 | assert mapping == [0, 1, 2, 3, 4, 5, 6] |