| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """int64 indexing in the block-verification rejection sampler kernels. |
| 4 | |
| 5 | With a GLM-scale vocab (~155k), logit_idx * vocab_stride exceeds int32 once |
| 6 | logit_idx >= ~13.8k (and req_state_idx * draft_stride_0 similarly), so the |
| 7 | block-verification kernels must promote indices to int64 before the stride |
| 8 | multiply. These tests place one request at a high logit/request-state index |
| 9 | and check its outputs match the same request run alone at index 0. |
| 10 | |
| 11 | Each case allocates ~5 GiB of GPU memory. |
| 12 | """ |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample |
| 19 | |
| 20 | VOCAB_SIZE = 155264 |
| 21 | NUM_SPECULATIVE_STEPS = 2 |
| 22 | LOGITS_PER_REQ = NUM_SPECULATIVE_STEPS + 1 |
| 23 | |
| 24 | |
| 25 | def _run( |
| 26 | target_logits: torch.Tensor, # [LOGITS_PER_REQ, V], the request under test |
| 27 | draft_logits: torch.Tensor, # [NUM_SPECULATIVE_STEPS, V] |
| 28 | draft_tokens: list[int], |
| 29 | num_reqs: int, |
| 30 | max_num_reqs: int, |
| 31 | roi_req: int, |
| 32 | req_state_idx: int, |
| 33 | device: torch.device, |
| 34 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 35 | """Run rejection_sample on a batch of uniform k=2 requests, with the |
| 36 | request under test at request index `roi_req` and draft-logits row |
| 37 | `req_state_idx`. Returns that request's (sampled, num_sampled).""" |
| 38 | num_logits = num_reqs * LOGITS_PER_REQ |
| 39 | roi_row = roi_req * LOGITS_PER_REQ |
| 40 | |
| 41 | full_target = torch.randn( |
| 42 | num_logits, VOCAB_SIZE, dtype=torch.bfloat16, device=device |
| 43 | ) |
| 44 | full_target[roi_row : roi_row + LOGITS_PER_REQ] = target_logits |
| 45 | full_draft = torch.randn( |
| 46 | max_num_reqs, |
| 47 | NUM_SPECULATIVE_STEPS, |
| 48 | VOCAB_SIZE, |
| 49 | dtype=torch.bfloat16, |
| 50 | device=device, |
| 51 | ) |
| 52 | full_draft[req_state_idx] = draft_logits |
| 53 | |
| 54 | # Draft token for step i of request r lives at row r*LOGITS_PER_REQ + i + 1. |
| 55 | draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device) |
| 56 | for i, tok in enumerate(draft_tokens): |
| 57 | draft_sampled[roi_row + i + 1] = tok |
| 58 | |
| 59 | cu_num_logits = torch.arange( |
| 60 | 0, num_logits + 1, LOGITS_PER_REQ, dtype=torch.int32, device=device |
| 61 | ) |
| 62 | # Filler requests all map to draft-logits row 0. |
| 63 | idx_mapping = torch.zeros(num_reqs, dtype=torch.int32, device=device) |
| 64 | idx_mapping[roi_req] = req_state_idx |
| 65 | expanded_idx_mapping = idx_mapping.repeat_interleave(LOGITS_PER_REQ) |
| 66 | expanded_local_pos = ( |
| 67 | torch.arange(LOGITS_PER_REQ, dtype=torch.int32, device=device) |
| 68 | .repeat(num_reqs) |
| 69 | .contiguous() |
| 70 | ) |
| 71 | |
| 72 | # Positions only need to match between the reference and big-batch runs |
| 73 | # for the request under test (they key the Gumbel noise and uniforms). |
| 74 | pos = torch.arange(num_logits, dtype=torch.int64, device=device) |
| 75 | pos[roi_row : roi_row + LOGITS_PER_REQ] = torch.arange( |
| 76 | 1000, 1000 + LOGITS_PER_REQ, dtype=torch.int64, device=device |
| 77 | ) |
| 78 | temperature = torch.ones(max_num_reqs, dtype=torch.float32, device=device) |
| 79 | seeds = torch.full((max_num_reqs,), 42, dtype=torch.int64, device=device) |
| 80 | |
| 81 | sampled, num_sampled = rejection_sample( |
| 82 | target_logits=full_target, |
| 83 | draft_logits=full_draft, |
| 84 | draft_sampled=draft_sampled, |
| 85 | cu_num_logits=cu_num_logits, |
| 86 | pos=pos, |
| 87 | idx_mapping=idx_mapping, |
| 88 | expanded_idx_mapping=expanded_idx_mapping, |
| 89 | expanded_local_pos=expanded_local_pos, |
| 90 | temperature=temperature, |
| 91 | seed=seeds, |
| 92 | num_speculative_steps=NUM_SPECULATIVE_STEPS, |
| 93 | use_block_verification=True, |
| 94 | ) |
| 95 | return sampled[roi_req].clone(), num_sampled[roi_req].clone() |
| 96 | |
| 97 | |
| 98 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") |
| 99 | @pytest.mark.parametrize( |
| 100 | # Overflow triggers when roi_req * LOGITS_PER_REQ * VOCAB_SIZE (target |
| 101 | # logits) or req_state_idx * NUM_SPECULATIVE_STEPS * VOCAB_SIZE (draft |
| 102 | # logits) exceeds 2**31. |
| 103 | ("num_reqs", "max_num_reqs", "roi_req", "req_state_idx"), |
| 104 | [ |
| 105 | (5500, 4, 5000, 2), # target-side: logit_idx 15000 * V > 2**31 |
| 106 | (22, 8192, 11, 8000), # draft-side: req_state_idx 8000 * 2V > 2**31 |
| 107 | ], |
| 108 | ) |
| 109 | def test_block_verification_i64_indexing( |
| 110 | num_reqs: int, max_num_reqs: int, roi_req: int, req_state_idx: int |
| 111 | ): |
| 112 | device = torch.device("cuda:0") |
| 113 | torch.manual_seed(0) |
| 114 | target_logits = torch.randn( |
| 115 | LOGITS_PER_REQ, VOCAB_SIZE, dtype=torch.bfloat16, device=device |
| 116 | ) |
| 117 | draft_logits = torch.randn( |
| 118 | NUM_SPECULATIVE_STEPS, VOCAB_SIZE, dtype=torch.bfloat16, device=device |
| 119 | ) |
| 120 | draft_tokens = [123, 45678] |
| 121 | |
| 122 | sampled_ref, num_sampled_ref = _run( |
| 123 | target_logits, |
| 124 | draft_logits, |
| 125 | draft_tokens, |
| 126 | num_reqs=1, |
| 127 | max_num_reqs=1, |
| 128 | roi_req=0, |
| 129 | req_state_idx=0, |
| 130 | device=device, |
| 131 | ) |
| 132 | sampled, num_sampled = _run( |
| 133 | target_logits, |
| 134 | draft_logits, |
| 135 | draft_tokens, |
| 136 | num_reqs=num_reqs, |
| 137 | max_num_reqs=max_num_reqs, |
| 138 | roi_req=roi_req, |
| 139 | req_state_idx=req_state_idx, |
| 140 | device=device, |
| 141 | ) |
| 142 | assert num_sampled.item() == num_sampled_ref.item() |
| 143 | n = num_sampled_ref.item() |
| 144 | assert torch.equal(sampled[:n], sampled_ref[:n]) |