from_numpy

vllm-project · Aug 5, 2024 · 4cdfa2e · 4cdfa2e
1 parent 90ecb1d
commit 4cdfa2e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/vllm/attention/backends/flash_attn.py b/vllm/attention/backends/flash_attn.py
@@ -310,7 +310,7 @@ def build(self, seq_lens: List[int], query_lens: List[int],
             for i, block_table in enumerate(self.block_tables):
                 if block_table:
                     input_block_tables[i, :len(block_table)] = block_table
-            block_tables = torch.tensor(input_block_tables).to(
+            block_tables = torch.from_numpy(input_block_tables).to(
                 device=device, non_blocking=True)
         else:
             block_tables = make_tensor_with_pad(

diff --git a/vllm/attention/backends/flashinfer.py b/vllm/attention/backends/flashinfer.py
@@ -357,7 +357,7 @@ def build(self, seq_lens: List[int], query_lens: List[int],
             for i, block_table in enumerate(self.block_tables):
                 if block_table:
                     input_block_tables[i, :len(block_table)] = block_table
-            block_tables = torch.tensor(input_block_tables).to(
+            block_tables = torch.from_numpy(input_block_tables).to(
                 device, non_blocking=True)
 
             last_paged_kv_indptr = self.paged_kv_indptr[-1]

diff --git a/vllm/attention/backends/utils.py b/vllm/attention/backends/utils.py
@@ -181,7 +181,7 @@ def build(self, seq_lens: List[int], query_lens: List[int],
             for i, block_table in enumerate(self.block_tables):
                 if block_table:
                     input_block_tables[i, :len(block_table)] = block_table
-            block_tables = torch.tensor(input_block_tables).to(
+            block_tables = torch.from_numpy(input_block_tables).to(
                 device, non_blocking=True)
         else:
             block_tables = make_tensor_with_pad(