Revert #7509 (#7887)

vllm-project · Aug 27, 2024 · 9606c71 · 9606c71
1 parent 64cc644
commit 9606c71
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/vllm/attention/backends/flashinfer.py b/vllm/attention/backends/flashinfer.py
@@ -113,8 +113,7 @@ def _get_decode_wrapper(self):
                 self.runner.parallel_config))
             num_kv_heads = self.runner.model_config.get_num_kv_heads(
                 self.runner.parallel_config)
-            use_tensor_cores = (num_qo_heads // num_kv_heads) not in \
-                (1, 2, 4, 8)
+            use_tensor_cores = num_qo_heads // num_kv_heads > 4
             self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
                 self._get_workspace_buffer(),
                 "NHD",
@@ -172,8 +171,7 @@ def graph_capture_get_metadata_for_batch(self, batch_size: int):
             self.runner.parallel_config))
         num_kv_heads = self.runner.model_config.get_num_kv_heads(
             self.runner.parallel_config)
-        use_tensor_cores = (num_qo_heads // num_kv_heads) not in \
-            (1, 2, 4, 8)
+        use_tensor_cores = num_qo_heads // num_kv_heads > 4
         self._graph_decode_wrapper = \
             CUDAGraphBatchDecodeWithPagedKVCacheWrapper(
             self._graph_decode_workspace_buffer, _indptr_buffer,