[Spec] Fold DFlash verified_id into the shared bonus_tokens relay channel (#29118)

python/sglang/srt/managers/overlap_utils.py

@@ -147,7 +147,6 @@ class FutureMap:
     def _lazy_init_forward_buf(self, draft_input: EagleDraftInput):
         self._forward_buf_initialized = True
 
-        self.need_verified_id = getattr(draft_input, "verified_id", None) is not None
         self.need_bonus_tokens = getattr(draft_input, "bonus_tokens", None) is not None
         self.need_topk = self.spec_algo.need_topk()
         self.need_hidden_states = (
@@ -155,22 +154,6 @@ class FutureMap:
             and getattr(draft_input, "hidden_states", None) is not None
         )
 
-        if self.need_verified_id:
-            verified_id0 = draft_input.verified_id[0]
-            self.verified_id_buf = (
-                torch.full(
-                    (self.req_pool_size, *verified_id0.shape),
-                    -1,
-                    dtype=verified_id0.dtype,
-                    device=self.device,
-                )
-                if _DEBUG_ASSERT
-                else torch.empty(
-                    (self.req_pool_size, *verified_id0.shape),
-                    dtype=verified_id0.dtype,
-                    device=self.device,
-                )
-            )
         if self.need_topk:
             topk_p0 = draft_input.topk_p[0]
             topk_index0 = draft_input.topk_index[0]
@@ -219,8 +202,6 @@ class FutureMap:
         # FIXME: indices = batch.req_pool_indices, pinned 2 iters via
         # record_batch_in_overlap; record_stream here is redundant.
         indices.record_stream(torch.get_device_module(self.device).current_stream())
-        if self.need_verified_id:
-            draft_input.verified_id = self.verified_id_buf[indices]
         if self.need_topk:
             hidden_states_buf = (
                 self.hidden_states_buf if self.need_hidden_states else None
@@ -248,10 +229,6 @@ class FutureMap:
         if self.need_hidden_states and not self.need_topk:
             draft_input.hidden_states = self.hidden_states_buf[indices]
         if _DEBUG_ASSERT:
-            if self.need_verified_id:
-                _assert_nonneg_and_invalidate(
-                    draft_input.verified_id, self.verified_id_buf, indices
-                )
             if self.need_bonus_tokens:
                 _assert_nonneg_and_invalidate(
                     draft_input.bonus_tokens, self.output_tokens_buf, indices
@@ -345,10 +322,6 @@ class FutureMap:
         draft_input: EagleDraftInput = payload
         if not self._forward_buf_initialized:
             self._lazy_init_forward_buf(draft_input)
-        if self.need_verified_id:
-            self.verified_id_buf[indices] = draft_input.verified_id.to(
-                self.verified_id_buf.dtype
-            )
         if self.need_bonus_tokens:
             self.output_tokens_buf[indices] = draft_input.bonus_tokens.to(
                 self.output_tokens_buf.dtype

python/sglang/srt/speculative/dflash_info_v2.py

@@ -41,11 +41,11 @@ class DFlashDraftInputV2(SpecInput):
     """Draft-side state carried across overlap iterations (spec-v2)."""
 
     # Legacy Eagle-shaped fields kept only for dataclass compatibility. DFLASH
-    # overlap carries new_seq_lens / verified_id directly in the common
+    # overlap carries new_seq_lens / bonus_tokens directly in the common
     # no-shape-change path; FutureMap remains the fallback for filter/merge.
     topk_p: torch.Tensor
     topk_index: torch.Tensor
-    verified_id: torch.Tensor
+    bonus_tokens: torch.Tensor
     new_seq_lens: torch.Tensor
     hidden_states: torch.Tensor
     verify_done: Optional[torch.cuda.Event] = None
@@ -122,7 +122,7 @@ class DFlashDraftInputV2(SpecInput):
         return cls(
             topk_p=torch.empty((0, 0), device=device, dtype=torch.float32),
             topk_index=torch.empty((0, 0), device=device, dtype=torch.int64),
-            verified_id=torch.empty((0,), device=device, dtype=torch.int32),
+            bonus_tokens=torch.empty((0,), device=device, dtype=torch.int64),
             new_seq_lens=torch.empty((0,), device=device, dtype=torch.int64),
             hidden_states=torch.empty((0, 0), device=device, dtype=torch.float16),
             verify_done=None,
@@ -298,7 +298,7 @@ class DFlashDraftInputV2(SpecInput):
 
         self.topk_p = self.topk_p[new_indices]
         self.topk_index = self.topk_index[new_indices]
-        self.verified_id = self.verified_id[new_indices]
+        self.bonus_tokens = self.bonus_tokens[new_indices]
         self.new_seq_lens = self.new_seq_lens[new_indices]
         self.hidden_states = self.hidden_states[new_indices]
 
@@ -341,7 +341,9 @@ class DFlashDraftInputV2(SpecInput):
 
         self.topk_p = torch.cat([self.topk_p, spec_info.topk_p], dim=0)
         self.topk_index = torch.cat([self.topk_index, spec_info.topk_index], dim=0)
-        self.verified_id = torch.cat([self.verified_id, spec_info.verified_id], dim=0)
+        self.bonus_tokens = torch.cat(
+            [self.bonus_tokens, spec_info.bonus_tokens], dim=0
+        )
         self.new_seq_lens = torch.cat(
             [self.new_seq_lens, spec_info.new_seq_lens], dim=0
         )

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -1160,17 +1160,17 @@ class DFlashWorkerV2(BaseSpecWorker):
     def _make_next_draft_input_prefill(
         self,
         *,
-        verified_id: torch.Tensor,
+        bonus_tokens: torch.Tensor,
         seq_lens: torch.Tensor,
         verify_done: Optional[torch.cuda.Event] = None,
         cur_allocated_seq_lens_cpu: Optional[torch.Tensor] = None,
     ) -> DFlashDraftInputV2:
         bs = int(seq_lens.numel())
-        device = verified_id.device
+        device = bonus_tokens.device
         return DFlashDraftInputV2(
             topk_p=torch.empty((bs, 0), device=device, dtype=torch.float32),
             topk_index=torch.empty((bs, 0), device=device, dtype=torch.int64),
-            verified_id=verified_id.to(dtype=torch.int32),
+            bonus_tokens=bonus_tokens.to(dtype=torch.int64),
             new_seq_lens=seq_lens.to(dtype=torch.int64),
             hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16),
             verify_done=verify_done,
@@ -1180,17 +1180,17 @@ class DFlashWorkerV2(BaseSpecWorker):
     def _make_next_draft_input_decode(
         self,
         *,
-        verified_id: torch.Tensor,
+        bonus_tokens: torch.Tensor,
         new_seq_lens: torch.Tensor,
         verify_done: Optional[torch.cuda.Event] = None,
         cur_allocated_seq_lens_cpu: Optional[torch.Tensor] = None,
     ) -> DFlashDraftInputV2:
         bs = int(new_seq_lens.numel())
-        device = verified_id.device
+        device = bonus_tokens.device
         return DFlashDraftInputV2(
             topk_p=torch.empty((bs, 0), device=device, dtype=torch.float32),
             topk_index=torch.empty((bs, 0), device=device, dtype=torch.int64),
-            verified_id=verified_id.to(dtype=torch.int32),
+            bonus_tokens=bonus_tokens.to(dtype=torch.int64),
             new_seq_lens=new_seq_lens.to(dtype=torch.int64),
             hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16),
             verify_done=verify_done,
@@ -1271,7 +1271,7 @@ class DFlashWorkerV2(BaseSpecWorker):
             logits_output.hidden_states = None
 
             batch_output.next_draft_input = self._make_next_draft_input_prefill(
-                verified_id=next_token_ids,
+                bonus_tokens=next_token_ids,
                 seq_lens=model_worker_batch.seq_lens,
                 cur_allocated_seq_lens_cpu=model_worker_batch.seq_lens_cpu,
             )
@@ -1296,7 +1296,7 @@ class DFlashWorkerV2(BaseSpecWorker):
             empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device)
             empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device)
             next_draft_input = self._make_next_draft_input_decode(
-                verified_id=torch.empty((0,), device=self.device, dtype=torch.int32),
+                bonus_tokens=torch.empty((0,), device=self.device, dtype=torch.int64),
                 new_seq_lens=torch.empty((0,), device=self.device, dtype=torch.int64),
             )
             if on_publish is not None:
@@ -1348,7 +1348,7 @@ class DFlashWorkerV2(BaseSpecWorker):
         if self._use_triton_prepare_block:
             try:
                 _prepare_dflash_draft_block_unchecked(
-                    verified_id=draft_input.verified_id.view(-1),
+                    bonus_tokens=draft_input.bonus_tokens.view(-1),
                     prefix_lens=prefix_lens.view(-1),
                     req_pool_indices=model_worker_batch.req_pool_indices.view(-1),
                     req_to_token=self.model_runner.req_to_token_pool.req_to_token,
@@ -1364,7 +1364,7 @@ class DFlashWorkerV2(BaseSpecWorker):
                     e,
                 )
                 block_ids.fill_(int(self._mask_token_id))
-                block_ids[:, 0].copy_(draft_input.verified_id)
+                block_ids[:, 0].copy_(draft_input.bonus_tokens)
                 torch.add(
                     prefix_lens.unsqueeze(1),
                     self._block_pos_offsets,
@@ -1383,7 +1383,7 @@ class DFlashWorkerV2(BaseSpecWorker):
                 verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size))
         else:
             block_ids.fill_(int(self._mask_token_id))
-            block_ids[:, 0].copy_(draft_input.verified_id)
+            block_ids[:, 0].copy_(draft_input.bonus_tokens)
             torch.add(
                 prefix_lens.unsqueeze(1),
                 self._block_pos_offsets,
@@ -1671,7 +1671,7 @@ class DFlashWorkerV2(BaseSpecWorker):
         logits_output.hidden_states = None
 
         next_draft_input = self._make_next_draft_input_decode(
-            verified_id=bonus,
+            bonus_tokens=bonus,
             new_seq_lens=new_seq_lens,
             cur_allocated_seq_lens_cpu=draft_input.reserved_seq_lens_cpu,
         )

python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py

@@ -5,14 +5,14 @@ import triton.language as tl
 
 @triton.jit
 def _prepare_dflash_draft_block_contig_kernel(
-    verified_id_ptr,
+    bonus_tokens_ptr,
     prefix_lens_ptr,
     req_pool_indices_ptr,
     req_to_token_ptr,
     block_ids_out_ptr,
     positions_out_ptr,
     cache_loc_out_ptr,
-    verified_id_stride,
+    bonus_tokens_stride,
     prefix_lens_stride,
     req_pool_indices_stride,
     req_to_token_row_stride,
@@ -30,7 +30,7 @@ def _prepare_dflash_draft_block_contig_kernel(
 
     prefix_len = tl.load(prefix_lens_ptr + row * prefix_lens_stride)
     req_idx = tl.load(req_pool_indices_ptr + row * req_pool_indices_stride)
-    verified_id = tl.load(verified_id_ptr + row * verified_id_stride)
+    bonus_token = tl.load(bonus_tokens_ptr + row * bonus_tokens_stride)
 
     logical_pos = prefix_len.to(tl.int64) + cols
     valid = row_mask & (logical_pos < req_to_token_width)
@@ -38,7 +38,7 @@ def _prepare_dflash_draft_block_contig_kernel(
     slot_ids = tl.load(req_row_ptr + logical_pos, mask=valid, other=0)
 
     block_ids = tl.full((BLOCK_SIZE,), mask_token_id, tl.int64)
-    block_ids = tl.where(cols == 0, verified_id.to(tl.int64), block_ids)
+    block_ids = tl.where(cols == 0, bonus_token.to(tl.int64), block_ids)
     tl.store(
         block_ids_out_ptr + row * block_ids_row_stride + cols, block_ids, mask=row_mask
     )
@@ -69,7 +69,7 @@ def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
 
 
 def _prepare_dflash_draft_block_unchecked(
-    verified_id: torch.Tensor,
+    bonus_tokens: torch.Tensor,
     prefix_lens: torch.Tensor,
     req_pool_indices: torch.Tensor,
     req_to_token: torch.Tensor,
@@ -78,7 +78,7 @@ def _prepare_dflash_draft_block_unchecked(
     cache_loc_out: torch.Tensor,
     mask_token_id: int,
 ) -> None:
-    batch_size = int(verified_id.numel())
+    batch_size = int(bonus_tokens.numel())
     if batch_size == 0:
         return
 
@@ -101,14 +101,14 @@ def _prepare_dflash_draft_block_unchecked(
     block = triton.next_power_of_2(block_size)
     num_warps = _pick_num_warps(block)
     _prepare_dflash_draft_block_contig_kernel[(batch_size,)](
-        verified_id,
+        bonus_tokens,
         prefix_lens,
         req_pool_indices,
         req_to_token,
         block_ids_out,
         positions_out,
         cache_loc_out,
-        verified_id.stride(0),
+        bonus_tokens.stride(0),
         prefix_lens.stride(0),
         req_pool_indices.stride(0),
         req_to_token.stride(0),