[Bugfix][Spec Decode][Structured Output] DSpark: fix the grammar bitmask mapping when the draft budget is zero (vllm-project#52436)

Signed-off-by: oops-oom <73481342@qq.com>
Co-authored-by: oops-oom <73481342@qq.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Wyett <wyettzeng@gmail.com>

Author: 3 people

tests/v1/spec_decode/test_adaptive_verification.py

@@ -9,6 +9,7 @@
 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]

vllm/v1/worker/gpu/model_runner.py

@@ -425,6 +425,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
                 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:

vllm/v1/worker/gpu/structured_outputs.py

@@ -10,13 +10,40 @@
 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,
         max_num_logits: int,
         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 @@ def __init__(
         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 @@ def apply_grammar_bitmask(
             )
 
         # 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):