fix: address Copilot review feedback — CPU tensor, redundant assignment, precompute max word idx

Author: windreamer

lmdeploy/pytorch/spec_decode/proposers/eagle3.py

@@ -34,6 +34,8 @@ def _init_bitmask_translate_constants(self):
         self._draft_words = draft_indices // 32
         self._draft_bits = draft_indices % 32
         self._n_draft_words = (draft_vocab_size + 31) // 32
+        # Precompute max word index (avoids GPU→CPU sync in _translate_bitmask)
+        self._max_d2t_word = self._d2t_words.max().item()
         # Cache device-specific constants; keyed by device.
         self._bitmask_cache: dict[torch.device, dict] = {}
 
@@ -66,7 +68,7 @@ def _translate_bitmask(self, target_bitmask: torch.Tensor) -> torch.Tensor:
         draft_words = c['draft_words']
         draft_bits = c['draft_bits']
 
-        max_word_idx = d2t_words.max().item()
+        max_word_idx = self._max_d2t_word
         if max_word_idx >= target_bitmask.size(1):
             raise ValueError(
                 f'd2t mapping references word index {max_word_idx} but target_bitmask '

lmdeploy/pytorch/spec_decode/spec_agent.py

@@ -497,7 +497,6 @@ def __compute_logprobs(raw_logprobs: torch.Tensor, token_ids: torch.LongTensor,
                 bonus_sampling_inputs,
                 logprobs_mode=self.misc_config.logprobs_mode,
             )
-            logits_processor.sampling_inputs = bonus_sampling_inputs
 
             next_token_ids = logits_processor.sampling(bonus_logits)  # [batch_size]
 
@@ -559,7 +558,8 @@ def __compute_logprobs(raw_logprobs: torch.Tensor, token_ids: torch.LongTensor,
                         _accept_spec_rejection_tokens,
                         guided_manager,
                         guided_processors,
-                        torch.zeros_like(next_token_ids),  # 0 rejected
+                        torch.zeros(next_token_ids.shape,
+                                    dtype=next_token_ids.dtype),  # 0 rejected, CPU
                         output_token_ids.cpu(),
                         cpu_next_token_ids,
                         0,  # num_spec_tokens=0, only bonus accepted