diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index fc6ef4a77..4b16610d3 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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 diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 9eaa467ff..aef81fc55 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/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 ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index c92a5d702..d8670e4b9 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py b/python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py index 4cfc8fe06..d14b0dd76 100644 --- a/python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py +++ b/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),