diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 5fc88de998..99708f50f8 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -9,6 +9,7 @@ from vllm.v1.worker.gpu.async_utils import StepTimingSample from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, ) +from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping def make_manager( @@ -168,3 +169,51 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits(): assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype # The prefill keeps its scheduled tokens; only drafts are dropped. assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) + + +def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): + # The scheduler sizes the grammar bitmask from the *scheduled* drafts + # (len(drafts) + 1 rows per request), but a zero budget rewrites + # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping + # from those rewritten offsets drops rows and trips the + # `num_masks == len(mapping)` assert in apply_grammar_bitmask. + manager = make_manager( + np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), + np.ones(64), + ) + manager.req_states.req_id_to_index["prefill"] = 2 + manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) + manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) + manager._max_total_logits = 2 # < 3 requests * 1 bonus token + + scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]} + manager.get_num_tokens( + {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens + ) + assert manager._batch_budget[2] == 0 + + req_ids = ["low", "high", "prefill"] + num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) + _, cu_num_logits_np = manager.compact_batch( + num_draft_tokens_per_req, + np.array([3, 3, 40], dtype=np.int32), + np.array([0, 3, 6, 7], dtype=np.int32), + ) + + mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens + mapping = _build_grammar_mapping( + req_ids, + req_ids, + cu_num_logits_np, + num_draft_tokens_per_req, + manager.num_bonus_tokens, + mask_stride, + ) + + num_bitmask_rows = sum( + len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids + ) + assert len(mapping) == num_bitmask_rows + # (request, position) keys, so the kernel can mask rows the compacted + # device layout no longer has room for. + assert mapping == [0, 1, 2, 3, 4, 5, 6] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index abbe6d5d3d..d68666cf4a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -425,6 +425,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): vocab_size=self.vocab_size, device=self.device, mask_stride=self.decode_query_len, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 221163ae36..8437b0aed9 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -10,6 +10,32 @@ from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch +def _build_grammar_mapping( + req_ids: list[str], + grammar_req_ids: list[str], + cu_num_logits_np: np.ndarray, + num_draft_tokens_per_req: np.ndarray | None, + num_bonus_tokens: int, + mask_stride: int, +) -> list[int]: + mapping: list[int] = [] + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} + for grammar_req_id in grammar_req_ids: + req_idx = req_id_to_idx[grammar_req_id] + if num_draft_tokens_per_req is None: + num_positions = int( + cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] + ) + else: + # Grammar masks follow the scheduled layout even when adaptive + # verification compacts the actual CPU logit offsets to bonus-only. + num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens + mapping.extend( + req_idx * mask_stride + position for position in range(num_positions) + ) + return mapping + + class StructuredOutputsWorker: def __init__( self, @@ -17,6 +43,7 @@ class StructuredOutputsWorker: vocab_size: int, device: torch.device, mask_stride: int, + num_bonus_tokens: int, ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device @@ -27,6 +54,7 @@ class StructuredOutputsWorker: self.device = device self.copy_stream = torch.cuda.Stream() self.mask_stride = mask_stride + self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -45,21 +73,17 @@ class StructuredOutputsWorker: ) # Construct bitmask -> logits mapping - mapping: list[int] = [] - req_ids = input_batch.req_ids - cu_num_logits = input_batch.cu_num_logits_np.tolist() - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - logits_start_idx = cu_num_logits[req_idx] - logits_end_idx = cu_num_logits[req_idx + 1] - # Key by (request, position) rather than absolute logit index: - # adaptive verification finalizes per-request logit offsets on - # device, so the kernel resolves them from the GPU cu_num_logits. - mapping.extend( - req_idx * self.mask_stride + position - for position in range(logits_end_idx - logits_start_idx) - ) + # Key by (request, position) rather than absolute logit index: + # adaptive verification finalizes per-request logit offsets on + # device, so the kernel resolves them from the GPU cu_num_logits. + mapping = _build_grammar_mapping( + input_batch.req_ids, + grammar_req_ids, + input_batch.cu_num_logits_np, + input_batch.num_draft_tokens_per_req, + self.num_bonus_tokens, + self.mask_stride, + ) # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream):