Merge branch 'v0.6.2-dev_wm' into 'v0.6.2-dev'

[fix]修复test_spec_decode_worker中的错误

See merge request dcutoolkit/deeplearing/vllm!34

tests/worker/test_model_input.py

@@ -60,6 +60,16 @@ class MockAttentionBackend(AttentionBackend):
     ) -> None:
         pass
 
+    @staticmethod
+    def move_cache(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+        kv_cache_dtype: str, 
+        num_kv_heads: int, 
+        head_size: int,
+    ) -> None:
+        NotImplementedError
+
 
 def test_model_runner_input():
     sampling_metadata = SamplingMetadata(

vllm/attention/backends/flash_attn.py

@@ -221,6 +221,16 @@ class FlashAttentionBackend(AttentionBackend):
         value_caches = [kv_cache[1] for kv_cache in kv_caches]
         ops.copy_blocks(key_caches, value_caches, src_to_dists)
 
+    @staticmethod
+    def move_cache(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+        kv_cache_dtype: str, 
+        num_kv_heads: int, 
+        head_size: int,
+    ) -> None:
+        NotImplementedError
+
 
 @dataclass
 class FlashAttentionMetadata(AttentionMetadata):

vllm/spec_decode/spec_decode_worker.py

@@ -550,14 +550,15 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
                     hidden_states, execute_model_req.seq_group_metadata_list)
                 
         # Store logits from target model execution.
-        logits = sampler_output.logits
-        if logits is not None:
-            if self.previous_logits is None:
-                self.previous_logits = Logits(
-                    logits, execute_model_req.seq_group_metadata_list)
-            else:
-                self.previous_logits.update(
-                    logits, execute_model_req.seq_group_metadata_list)
+        if self.tree_style_spec_decoding:
+            logits = sampler_output.logits
+            if logits is not None:
+                if self.previous_logits is None:
+                    self.previous_logits = Logits(
+                        logits, execute_model_req.seq_group_metadata_list)
+                else:
+                    self.previous_logits.update(
+                        logits, execute_model_req.seq_group_metadata_list)
 
         if not skip_proposer:
             # We prepare the prefill hidden states here so that there no